# noqa: PLC0302
"""Misc collection of tests for nodes in this extension"""

import math
import os
from functools import partial
from math import isnan
from typing import Callable, Set

import carb
import carb.settings
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.commands
import omni.kit.stage_templates
import omni.kit.test
import omni.timeline
import omni.usd
import usdrt
from omni.graph.core import ThreadsafetyTestUtils
from pxr import Gf, OmniGraphSchemaTools, Sdf, Usd, UsdGeom, UsdShade


# ======================================================================
class TestNodes(ogts.OmniGraphTestCase):
    """Unit tests nodes in this extension"""

    PERSISTENT_SETTINGS_PREFIX = "/persistent"
    TEST_GRAPH_PATH = "/World/TestGraph"

    async def setUp(self):
        """Set up  test environment, to be torn down when done"""
        await super().setUp()

        # Ensure that Xform prims are created with the full set of xformOp attributes with consistent behavior.
        settings = carb.settings.get_settings()
        settings.set(self.PERSISTENT_SETTINGS_PREFIX + "/app/primCreation/PrimCreationWithDefaultXformOps", True)
        settings.set(
            self.PERSISTENT_SETTINGS_PREFIX + "/app/primCreation/DefaultXformOpType", "Scale, Rotate, Translate"
        )
        settings.set(self.PERSISTENT_SETTINGS_PREFIX + "/app/primCreation/DefaultRotationOrder", "XYZ")
        settings.set(
            self.PERSISTENT_SETTINGS_PREFIX + "/app/primCreation/DefaultXformOpOrder",
            "xformOp:translate, xformOp:rotate, xformOp:scale",
        )
        settings.set(self.PERSISTENT_SETTINGS_PREFIX + "/app/primCreation/DefaultXformOpPrecision", "Double")

    # ----------------------------------------------------------------------
    async def test_getprimrelationship(self):
        """Test GetPrimRelationship node"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (get_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("Get", "omni.graph.nodes.GetPrimRelationship"),
                keys.CREATE_PRIMS: [("RelHolder", {}), ("TestPrimA", {}), ("TestPrimB", {})],
                keys.SET_VALUES: [("Get.inputs:name", "test_rel"), ("Get.inputs:usePath", False)],
            },
        )
        await controller.evaluate(graph)
        prim = stage.GetPrimAtPath("/RelHolder")
        rel = prim.CreateRelationship("test_rel")
        for prim_path in ("/TestPrimA", "/TestPrimB"):
            omni.kit.commands.execute("AddRelationshipTarget", relationship=rel, target=Sdf.Path(prim_path))

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Get.inputs:prim"),
            target=prim.GetPath(),
        )
        await controller.evaluate(graph)
        rel_paths = og.Controller.get(controller.attribute("outputs:paths", get_node))
        self.assertEqual(rel_paths, rel.GetTargets())

        controller.edit(
            self.TEST_GRAPH_PATH,
            {keys.SET_VALUES: [("Get.inputs:path", "/RelHolder"), ("Get.inputs:usePath", True)]},
        )
        await controller.evaluate(graph)
        rel_paths = og.Controller.get(controller.attribute("outputs:paths", get_node))
        self.assertEqual(rel_paths, rel.GetTargets())

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_get_prim_path(self, test_instance_id: int = 0):
        """Test GetPrimPath node"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)

        prim_path = "/World/TestPrim"
        ThreadsafetyTestUtils.add_to_threading_cache(test_instance_id, stage.DefinePrim(prim_path))

        controller = og.Controller()
        keys = og.Controller.Keys
        (_, (get_node,), _, _) = controller.edit(
            graph_path,
            {keys.CREATE_NODES: ("GetPrimPath", "omni.graph.nodes.GetPrimPath")},
        )

        rel = stage.GetPropertyAtPath(f"{graph_path}/GetPrimPath.inputs:prim")
        omni.kit.commands.execute("AddRelationshipTarget", relationship=rel, target=prim_path)

        yield ThreadsafetyTestUtils.EVALUATION_ALL_GRAPHS
        rel_paths = og.Controller.get(controller.attribute("outputs:primPath", get_node))
        self.assertEqual(rel_paths, prim_path)

        ThreadsafetyTestUtils.single_evaluation_last_test_instance(
            test_instance_id, lambda: stage.RemovePrim(prim_path)
        )

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_get_prim_paths(self, test_instance_id: int = 0):
        """Test GetPrimPaths node"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)

        prim_paths = ["/World/TestPrim1", "/World/TestPrim2"]
        ThreadsafetyTestUtils.add_to_threading_cache(test_instance_id, stage.DefinePrim(prim_paths[0]))
        ThreadsafetyTestUtils.add_to_threading_cache(test_instance_id, stage.DefinePrim(prim_paths[1]))

        controller = og.Controller()
        keys = og.Controller.Keys

        (_, (get_node,), _, _) = controller.edit(
            graph_path,
            {keys.CREATE_NODES: ("GetPrimPaths", "omni.graph.nodes.GetPrimPaths")},
        )

        rel = stage.GetPropertyAtPath(f"{graph_path}/GetPrimPaths.inputs:prims")
        omni.kit.commands.execute("AddRelationshipTarget", relationship=rel, target=prim_paths[0])
        omni.kit.commands.execute("AddRelationshipTarget", relationship=rel, target=prim_paths[1])

        yield ThreadsafetyTestUtils.EVALUATION_ALL_GRAPHS
        rel_paths = og.Controller.get(controller.attribute("outputs:primPaths", get_node))
        self.assertEqual(rel_paths, prim_paths)

        ThreadsafetyTestUtils.single_evaluation_last_test_instance(
            test_instance_id, lambda: stage.RemovePrim(prim_paths[0])
        )
        ThreadsafetyTestUtils.single_evaluation_last_test_instance(
            test_instance_id, lambda: stage.RemovePrim(prim_paths[1])
        )

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_get_prims_at_path(self, test_instance_id: int = 0):
        """Test GetPrimAtPath node"""
        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)

        controller = og.Controller()
        prim_path = "/World/foo/bar"
        keys = og.Controller.Keys
        (_, (get_prim_at_path, get_prims_at_path, _, _), _, _) = controller.edit(
            graph_path,
            {
                keys.CREATE_NODES: [
                    ("GetPrimAtPath", "omni.graph.nodes.GetPrimsAtPath"),
                    ("GetPrimsAtPath", "omni.graph.nodes.GetPrimsAtPath"),
                    ("MakeArray", "omni.graph.nodes.ConstructArray"),
                    ("ConstToken", "omni.graph.nodes.ConstantToken"),
                ],
                keys.CONNECT: [
                    ("ConstToken.inputs:value", "GetPrimAtPath.inputs:path"),
                    ("MakeArray.outputs:array", "GetPrimsAtPath.inputs:path"),
                ],
                keys.SET_VALUES: [
                    ("MakeArray.inputs:arraySize", 3),
                    ("MakeArray.inputs:arrayType", "token[]"),
                    ("MakeArray.inputs:input0", prim_path),
                    ("ConstToken.inputs:value", prim_path),
                ],
            },
        )
        yield ThreadsafetyTestUtils.EVALUATION_ALL_GRAPHS
        out_path = og.Controller.get(controller.attribute("outputs:prims", get_prim_at_path))
        self.assertEqual(out_path, [usdrt.Sdf.Path(prim_path)])
        out_paths = og.Controller.get(controller.attribute("outputs:prims", get_prims_at_path))
        self.assertEqual(out_paths, [usdrt.Sdf.Path(prim_path)] * 3)

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_constant_prims(self, test_instance_id: int = 0):
        """Test ConstantPrims node"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)

        prim_paths = ["/World/TestPrim1", "/World/TestPrim2"]
        for prim_path in prim_paths:
            ThreadsafetyTestUtils.add_to_threading_cache(test_instance_id, stage.DefinePrim(prim_path))

        controller = og.Controller()
        keys = og.Controller.Keys
        (_, (_const_prims, prim_paths_node), _, _) = controller.edit(
            graph_path,
            {
                keys.CREATE_NODES: [
                    ("ConstantPrims", "omni.graph.nodes.ConstantPrims"),
                    ("GetPrimPath", "omni.graph.nodes.GetPrimPaths"),
                ],
                keys.CONNECT: [("ConstantPrims.inputs:value", "GetPrimPath.inputs:prims")],
            },
        )

        rel = stage.GetPropertyAtPath(f"{graph_path}/ConstantPrims.inputs:value")
        for prim_path in prim_paths:
            omni.kit.commands.execute("AddRelationshipTarget", relationship=rel, target=prim_path)

        yield ThreadsafetyTestUtils.EVALUATION_ALL_GRAPHS

        out_paths = og.Controller.get(controller.attribute("outputs:primPaths", prim_paths_node))
        self.assertEqual(out_paths, prim_paths)

        for prim_path in prim_paths:
            ThreadsafetyTestUtils.single_evaluation_last_test_instance(
                test_instance_id, partial(stage.RemovePrim, prim_path)
            )

    # ----------------------------------------------------------------------
    async def test_constant_prims_loads_from_file(self):
        """
        Validation that connections are maintained when loading a graph from a file which uses ConstantPrims -
        which uses a target input attribute as its output
        """
        await ogts.load_test_file("TestConstantPrims.usda", use_caller_subdirectory=True)

        self.assertTrue(og.Controller.graph("/World/PushGraph").is_valid())
        self.assertTrue(og.Controller.node("/World/PushGraph/constant_prims").is_valid())
        self.assertTrue(og.Controller.node("/World/PushGraph/get_prim_paths").is_valid())
        attr_dst = og.Controller.attribute("/World/PushGraph/get_prim_paths.inputs:prims")
        attr_src = og.Controller.attribute("/World/PushGraph/constant_prims.inputs:value")
        self.assertTrue(attr_dst.is_valid())
        self.assertTrue(attr_src.is_valid())
        self.assertTrue(attr_dst.is_connected(attr_src))
        self.assertEqual(attr_dst.get_upstream_connection_count(), 1)
        self.assertEqual(attr_dst.get_upstream_connections()[0].get_path(), attr_src.get_path())
        actual_result = og.Controller.get(attr_dst)
        expected_result = [usdrt.Sdf.Path("/Environment"), usdrt.Sdf.Path("/Environment/defaultLight")]
        self.assertEqual(expected_result, actual_result)

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_get_parent_prims(self, test_instance_id: int = 0):
        """Test ConstantPrims node"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)

        prim_paths = ["/World/Test", "/World/Test/Test"]
        prim_parents = ["/World", "/World/Test"]
        for prim_path in prim_paths:
            ThreadsafetyTestUtils.add_to_threading_cache(test_instance_id, stage.DefinePrim(prim_path))

        controller = og.Controller()
        keys = og.Controller.Keys
        (_, (get_parent_prims,), _, _) = controller.edit(
            graph_path, {keys.CREATE_NODES: ("GetParentPrims", "omni.graph.nodes.GetParentPrims")}
        )

        rel = stage.GetPropertyAtPath(f"{graph_path}/GetParentPrims.inputs:prims")
        for prim_path in prim_paths:
            omni.kit.commands.execute("AddRelationshipTarget", relationship=rel, target=prim_path)

        yield ThreadsafetyTestUtils.EVALUATION_ALL_GRAPHS
        out_paths = og.Controller.get(controller.attribute("outputs:parentPrims", get_parent_prims))
        self.assertEqual(out_paths, [usdrt.Sdf.Path(p) for p in prim_parents])

        for prim_path in prim_paths:
            ThreadsafetyTestUtils.single_evaluation_last_test_instance(
                test_instance_id, partial(stage.RemovePrim, prim_path)
            )

    # ----------------------------------------------------------------------
    async def test_prim_relationship_load(self):
        """Test that a prim relationship when loading does not create OgnPrim node"""

        async def do_test():
            await ogts.load_test_file("TestPrimRelationshipLoad.usda", use_caller_subdirectory=True)
            # Cube and Capsule are driving Sphere, Cylinder, Cone by various methods
            usd_context = omni.usd.get_context()
            stage = usd_context.get_stage()
            graph = og.get_graph_by_path("/World/PushGraph")
            controller = og.Controller()
            nodes = graph.get_nodes()
            node_paths = [n.get_prim_path() for n in nodes]

            # These prims are only connected by relationship, so should not show up in OG
            self.assertNotIn("/World/Cube", node_paths)
            self.assertNotIn("/World/Sphere", node_paths)
            self.assertNotIn("/World/Cylinder", node_paths)

            self.assertNotIn("/World/Cone", node_paths)

            # Sanity check that the rotate connections are working
            for p in ("Sphere", "Capsule", "Cylinder", "Cone"):
                rotate_xyz = stage.GetPrimAtPath("/World/" + p).GetAttribute("xformOp:rotateXYZ").Get()
                self.assertEqual(rotate_xyz[0], 0)
                self.assertEqual(rotate_xyz[1], 0)

            stage.GetPrimAtPath("/World/Cube").GetAttribute("xformOp:rotateXYZ").Set((100, 0, 0))
            stage.GetPrimAtPath("/World/Capsule").GetAttribute("xformOp:rotateXYZ").Set((100, 0, 0))
            await controller.evaluate(graph)
            for p in ("Sphere", "Cylinder", "Cone"):
                rotate_xyz = stage.GetPrimAtPath("/World/" + p).GetAttribute("xformOp:rotateXYZ").Get()
                if p == "Cone":
                    # /World/Cone is not loaded into Flatcache, so we expect this part of the
                    # graph to be non-functional (Cone will not change)
                    self.assertEqual(rotate_xyz[0], 0)
                self.assertEqual(rotate_xyz[1], 0)

        await do_test()

    # ----------------------------------------------------------------------
    @staticmethod
    def _get_expected_read_prims_property_names(prim) -> Set[str]:
        properties = set(prim.GetAuthoredPropertyNames())
        properties.update({"worldMatrix", "sourcePrimPath", "sourcePrimType"})
        return properties

    # ----------------------------------------------------------------------
    @staticmethod
    def _get_expected_read_prims_v2_property_names(prim) -> Set[str]:
        properties = set(prim.GetPropertyNames())
        properties.remove("proxyPrim")  # remove "proxyPrim" which is skipped unless it has targets
        properties.update({"worldMatrix", "sourcePrimPath", "sourcePrimType"})
        return properties

    # ----------------------------------------------------------------------
    async def test_read_prims_write_prim(self):
        """Test omni.graph.nodes.ReadPrims and WritePrim"""
        await self._read_prims_write_prim(
            "omni.graph.nodes.ReadPrims",
            False,
            compute_expected_property_names=self._get_expected_read_prims_property_names,
        )

    # ----------------------------------------------------------------------
    async def test_read_prims_write_prim_with_target(self):
        await self._read_prims_write_prim(
            "omni.graph.nodes.ReadPrims",
            True,
            compute_expected_property_names=self._get_expected_read_prims_property_names,
        )

    # ----------------------------------------------------------------------
    async def test_read_prims_v2_write_prim(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrim"""
        await self._read_prims_write_prim(
            "omni.graph.nodes.ReadPrimsV2",
            False,
            compute_expected_property_names=self._get_expected_read_prims_v2_property_names,
        )

    # ----------------------------------------------------------------------
    async def test_read_prims_v2_write_prim_with_target(self):
        await self._read_prims_write_prim(
            "omni.graph.nodes.ReadPrimsV2",
            True,
            compute_expected_property_names=self._get_expected_read_prims_v2_property_names,
        )

    # ----------------------------------------------------------------------
    async def test_read_prims_bundle_write_prim(self):
        """Test omni.graph.nodes.ReadPrimsBundle and WritePrim"""
        await self._read_prims_write_prim(
            "omni.graph.nodes.ReadPrimsBundle",
            False,
            compute_expected_property_names=self._get_expected_read_prims_property_names,
        )

    # ----------------------------------------------------------------------
    async def test_read_prims_bundle_write_prim_with_target(self):
        await self._read_prims_write_prim(
            "omni.graph.nodes.ReadPrimsBundle",
            True,
            compute_expected_property_names=self._get_expected_read_prims_property_names,
        )

    # ----------------------------------------------------------------------
    async def _read_prims_write_prim(
        self, read_prims_type, use_target_inputs, compute_expected_property_names: Callable
    ):
        """
        use_target_inputs will use the prim target input for the ExtractPrim nodes rather than the prim path input
        """
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        cube_prim = ogts.create_cube(stage, "Cube", (1, 1, 1))

        (graph, (read_node, write_node, _, _, extract_bundle_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", read_prims_type),
                    ("Write", "omni.graph.nodes.WritePrim"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("ExtractPrim", "omni.graph.nodes.ExtractPrim"),
                    ("ExtractBundle", "omni.graph.nodes.ExtractBundle"),
                ],
                keys.SET_VALUES: [
                    ("ExtractPrim.inputs:primPath", "/Cube"),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "ExtractPrim.inputs:prims"),
                    ("ExtractPrim.outputs_primBundle", "ExtractBundle.inputs:bundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            target=cube_prim.GetPath(),
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prim"),
            target=cube_prim.GetPath(),
        )
        await controller.evaluate(graph)

        if use_target_inputs:
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetRelationshipAtPath(f"{self.TEST_GRAPH_PATH}/ExtractPrim.inputs:prim"),
                target=cube_prim.GetPath(),
            )
            await controller.evaluate(graph)

        factory = og.IBundleFactory.create()

        # Check MPiB output for Read
        bundle = graph.get_default_graph_context().get_output_bundle(read_node, "outputs_primsBundle")
        output_bundle2 = factory.get_bundle(graph.get_default_graph_context(), bundle)
        self.assertEqual(output_bundle2.get_child_bundle_count(), 1)

        # Attribute Count:
        # 2 - node: type, typeVersion,
        # 3 - inputs: bundle
        # 2 - outputs: passThrough,
        n_static_attribs_extract_bundle = 4

        # Attribute Count:
        # 2 - node: type, typeVersion,
        # 3 - inputs: prim, execIn, usdWriteBack
        # 1 - outputs: execOut,
        n_static_attribs_write = 6

        extract_bundle_name = "outputs_passThrough"

        attribs = extract_bundle_node.get_attributes()
        found_size_attrib = False
        for attrib in attribs:
            if attrib.get_name() == "outputs:size" and attrib.get_resolved_type().base_type == og.BaseDataType.DOUBLE:
                found_size_attrib = True
        self.assertTrue(found_size_attrib)

        expected_property_names = compute_expected_property_names(cube_prim)

        bundle = graph.get_default_graph_context().get_output_bundle(extract_bundle_node, extract_bundle_name)

        attribute_names = set(bundle.get_attribute_names())
        self.assertEqual(attribute_names, expected_property_names)

        # Check we have the expected bundle attributes
        self.assertIn("size", attribute_names)
        self.assertIn("sourcePrimPath", attribute_names)
        self.assertIn("sourcePrimType", attribute_names)
        self.assertIn("worldMatrix", attribute_names)
        self.assertIn("primvars:displayColor", attribute_names)

        attribute_datas = bundle.get_attribute_data(False)
        self.assertEqual(
            next((a for a in attribute_datas if a.get_name() == "sourcePrimPath")).get(),
            cube_prim.GetPath().pathString,
        )

        # Check we have the expected dynamic attrib on WritePrim
        attribs = write_node.get_attributes()
        found_size_attrib = False
        for attrib in attribs:
            if attrib.get_name() == "inputs:size" and attrib.get_resolved_type().base_type == og.BaseDataType.DOUBLE:
                found_size_attrib = True
        self.assertTrue(found_size_attrib)

        # check that evaluations propagate in a read/write cycle as expected
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CONNECT: [
                    ("ExtractBundle.outputs:size", "Add.inputs:a"),
                    ("ExtractBundle.outputs:size", "Add.inputs:b"),
                    ("Add.outputs:sum", "Write.inputs:size"),
                ]
            },
        )

        await controller.evaluate(graph)
        attr = controller.attribute("outputs:size", extract_bundle_node)
        # it is now 2 on prim, but read node as not been executed after the write, so still has the old value
        self.assertEqual(og.Controller.get(attr), 1)
        await controller.evaluate(graph)
        # now it is 4 on the prim, and /Read as 2
        self.assertEqual(og.Controller.get(attr), 2)
        await controller.evaluate(graph)
        # now it is 8 on the prim, and /Read as 4
        self.assertEqual(og.Controller.get(attr), 4)

        # ReadPrim: Clear the inputs:prim attribute to trigger a reset of the ReadPrim node
        omni.kit.commands.execute(
            "RemoveRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            target=cube_prim.GetPath(),
        )
        await controller.evaluate(graph)

        # Verify that the stale dynamic and bundle attributes have been removed
        attribs = extract_bundle_node.get_attributes()
        self.assertEqual(len(attribs), n_static_attribs_extract_bundle)

        bundle = graph.get_default_graph_context().get_output_bundle(extract_bundle_node, extract_bundle_name)
        self.assertEqual(bundle.get_attribute_data_count(), 0)

        # WritePrim: Clear the inputs:prim attribute to trigger a reset of the WritePrim node
        omni.kit.commands.execute(
            "RemoveRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prim"),
            target=cube_prim.GetPath(),
        )
        await controller.evaluate(graph)

        # Verify that the stale dynamic and bundle attributes have been removed
        attribs = write_node.get_attributes()
        self.assertEqual(len(attribs), n_static_attribs_write)

        # Set it up again and re-verify things work
        omni.kit.undo.undo()  # Undo remove target from Write Prim

        # Following step is to undo disconnected and removed dynamic attributes inside of Extract Bundle.
        # When Extract Bundle's input change or recompute is called, we record
        omni.kit.undo.undo()  # Call to undo dynamic attribute removal and disconnection on read prim (compute)
        omni.kit.undo.undo()  # Call to undo dynamic attribute removal and disconnection on read prim (compute)
        omni.kit.undo.undo()  # Call to undo dynamic attribute removal and disconnection on read prim into bundle (compute)
        omni.kit.undo.undo()  # Call to undo dynamic attribute removal and disconnection on read prim into bundle (compute)

        await controller.evaluate(graph)

        # Check if MPiB is back to the output of ReadPrimIntoBundle
        bundle = graph.get_default_graph_context().get_output_bundle(read_node, "outputs_primsBundle")
        output_bundle2 = factory.get_bundle(graph.get_default_graph_context(), bundle)
        self.assertEqual(output_bundle2.get_child_bundle_count(), 1)

        # Verify we have the size attribute back and our bundle attribs
        attribs = write_node.get_attributes()

        found_size_attrib = False
        for attrib in attribs:
            if attrib.get_name() == "inputs:size" and attrib.get_resolved_type().base_type == og.BaseDataType.DOUBLE:
                found_size_attrib = True
        self.assertTrue(found_size_attrib)

        # Check output bundle attributes of ExtractBundle
        bundle = graph.get_default_graph_context().get_output_bundle(extract_bundle_node, extract_bundle_name)
        attribute_datas = bundle.get_attribute_data(False)
        attribute_names = [a.get_name() for a in attribute_datas]
        self.assertIn("size", attribute_names)
        self.assertIn("sourcePrimPath", attribute_names)
        self.assertIn("worldMatrix", attribute_names)
        self.assertIn("primvars:displayColor", attribute_names)

        attribs = extract_bundle_node.get_attributes()
        self.assertEqual(len(attribs) - n_static_attribs_extract_bundle, len(attribute_datas))

        # Reset the underlying prim attrib value and check the plumbing is still working
        cube_prim.GetAttribute("size").Set(1)
        await controller.evaluate(graph)
        attr = controller.attribute("outputs:size", extract_bundle_node)
        self.assertEqual(og.Controller.get(attr), 1)
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(attr), 2)
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(attr), 4)

    # ----------------------------------------------------------------------
    async def test_findprims(self):
        """Test omni.graph.nodes.FindPrims corner cases"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        cube_prim = ogts.create_cube(stage, "Cube1", (1, 1, 1))
        cube_prim2 = ogts.create_cube(stage, "Cube2", (1, 1, 1))
        cube_prim3 = ogts.create_cube(stage, "Cube3", (1, 1, 1))
        cube_prim3.CreateRelationship("test_rel").AddTarget(cube_prim.GetPrimPath())

        (graph, (read_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("FindPrims", "omni.graph.nodes.FindPrims"),
                keys.SET_VALUES: ("FindPrims.inputs:namePrefix", "Cube"),
            },
        )
        await controller.evaluate(graph)
        paths_attr = controller.attribute("outputs:primPaths", read_node)
        paths = og.Controller.get(paths_attr)
        self.assertListEqual(paths, [p.GetPrimPath() for p in (cube_prim, cube_prim2, cube_prim3)])
        # set cube to inactive and verify it doesn't show up in the list
        cube_prim.SetActive(False)
        await controller.evaluate(graph)
        paths = og.Controller.get(paths_attr)
        self.assertListEqual(paths, [p.GetPrimPath() for p in (cube_prim2, cube_prim3)])

        # Test the relationship requirement works
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: [
                    ("FindPrims.inputs:requiredRelationship", "test_rel"),
                    # Not the right target
                    ("FindPrims.inputs:requiredRelationshipTarget", cube_prim2.GetPrimPath().pathString),
                ]
            },
        )
        await controller.evaluate(graph)
        paths = og.Controller.get(paths_attr)
        self.assertTrue(len(paths) == 0)

        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: [
                    # The right target
                    ("FindPrims.inputs:requiredRelationshipTarget", cube_prim.GetPrimPath().pathString),
                ]
            },
        )
        await controller.evaluate(graph)
        paths = og.Controller.get(paths_attr)
        self.assertListEqual(paths, [p.GetPrimPath() for p in (cube_prim3,)])

    # ----------------------------------------------------------------------
    async def test_findprims_path_pattern(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        ogts.create_cube(stage, "Cube1", (1, 1, 1))
        ogts.create_cube(stage, "Cube2", (1, 1, 1))
        ogts.create_cube(stage, "Cube3", (1, 1, 1))
        ogts.create_cube(stage, "Cube44", (1, 1, 1))
        ogts.create_cube(stage, "AnotherCube1", (1, 1, 1))
        ogts.create_cube(stage, "AnotherCube2", (1, 1, 1))
        ogts.create_cube(stage, "AnotherCube3", (1, 1, 1))

        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (find_prims,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("FindPrims", "omni.graph.nodes.FindPrims"),
                keys.SET_VALUES: [
                    ("FindPrims.inputs:recursive", True),
                    ("FindPrims.inputs:pathPattern", "/Cube?"),
                ],
            },
        )
        await controller.evaluate(graph)

        outputs_prim_paths = find_prims.get_attribute("outputs:primPaths")
        value = outputs_prim_paths.get()
        self.assertEqual(len(value), 3)  # Cube1, Cube2, Cube3 but not Cube44
        self.assertTrue(all(item in ["/Cube1", "/Cube2", "/Cube3"] for item in value))

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: ("FindPrims.inputs:pathPattern", "/*Cube?"),
            },
        )
        await controller.evaluate(graph)
        value = outputs_prim_paths.get()
        self.assertEqual(len(value), 6)  # AnotherCube1,2,3 and Cube1,2,3, but not Cube44
        required_cubes = ["/Cube1", "/Cube2", "/Cube3", "/AnotherCube1", "/AnotherCube2", "/AnotherCube3"]
        self.assertTrue(all(item in required_cubes for item in value))

        # get all /Cube?, but exclude /Cube2
        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: ("FindPrims.inputs:pathPattern", "/Cube? ^/Cube2"),
            },
        )
        await controller.evaluate(graph)
        value = outputs_prim_paths.get()
        self.assertEqual(len(value), 2)  # Cube1,3, but not Cube2
        required_cubes = ["/Cube1", "/Cube3"]
        self.assertTrue(all(item in required_cubes for item in value))

        # get all /Cube? but exclude /Cube1 and /Cube2
        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: ("FindPrims.inputs:pathPattern", "/Cube? ^/Cube1 ^/Cube2"),
            },
        )
        await controller.evaluate(graph)
        value = outputs_prim_paths.get()
        self.assertEqual(len(value), 1)  # Cube3, but not Cube1,2
        required_cubes = ["/Cube3"]
        self.assertTrue(all(item in required_cubes for item in value))

        # get all Cubes, but exclude all cubes that end with single digit
        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: ("FindPrims.inputs:pathPattern", "*Cube* ^*Cube?"),
            },
        )
        await controller.evaluate(graph)
        value = outputs_prim_paths.get()
        self.assertEqual(len(value), 1)  # only Cube44
        required_cubes = ["/Cube44"]
        self.assertTrue(all(item in required_cubes for item in value))

        # get all Cubes, and exclude all cubes to produce empty set
        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: ("FindPrims.inputs:pathPattern", "*Cube* ^*Cube*"),
            },
        )
        await controller.evaluate(graph)
        value = outputs_prim_paths.get()
        self.assertEqual(len(value), 0)

    # ----------------------------------------------------------------------
    async def test_getprims_path_pattern(self):
        """Test the path pattern matching feature of the GetPrims node."""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_, get_prims_node), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("ReadPrimsBundle", "omni.graph.nodes.ReadPrimsBundle"),
                    ("GetPrims", "omni.graph.nodes.GetPrims"),
                ],
                keys.CONNECT: [("ReadPrimsBundle.outputs_primsBundle", "GetPrims.inputs:bundle")],
            },
        )

        cube_paths = {"/Cube1", "/Cube2", "/Cube3", "/Cube44", "/AnotherCube1", "/AnotherCube2", "/AnotherCube3"}
        for cube_path in cube_paths:
            cube_name = cube_path.split("/")[-1]
            ogts.create_cube(stage, cube_name, (1, 1, 1))
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/ReadPrimsBundle.inputs:prims"),
                target=cube_path,
            )

        await controller.evaluate(graph)

        # Dictionary from each path pattern to its expected prim paths
        path_pattern_dict = {
            "*": cube_paths,
            "": set(),
            "/Cube?": {"/Cube1", "/Cube2", "/Cube3"},
            "/*Cube?": {"/Cube1", "/Cube2", "/Cube3", "/AnotherCube1", "/AnotherCube2", "/AnotherCube3"},
            "/* ^/Another*": {"/Cube1", "/Cube2", "/Cube3", "/Cube44"},
        }

        bundle_factory = og.IBundleFactory.create()

        bundle = graph.get_default_graph_context().get_output_bundle(get_prims_node, "outputs_bundle")
        output_bundle2 = bundle_factory.get_bundle(graph.get_default_graph_context(), bundle)

        for path_pattern, expected_prim_paths in path_pattern_dict.items():
            # Test path patterns with the "inverse" option off (by default)
            (graph, _, _, _) = controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.SET_VALUES: [
                        ("GetPrims.inputs:pathPattern", path_pattern),
                        ("GetPrims.inputs:inverse", False),
                    ]
                },
            )
            await controller.evaluate(graph)

            child_bundle_paths = set()

            child_bundle_count = output_bundle2.get_child_bundle_count()
            for i in range(child_bundle_count):
                child_bundle = output_bundle2.get_child_bundle(i)
                attr = child_bundle.get_attribute_by_name("sourcePrimPath")
                child_bundle_paths.add(attr.get())

            self.assertEqual(child_bundle_paths, expected_prim_paths)

            # Test path patterns with the "inverse" option on
            (graph, _, _, _) = controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.SET_VALUES: [
                        ("GetPrims.inputs:inverse", True),
                    ]
                },
            )
            await controller.evaluate(graph)

            child_bundle_paths = set()

            child_bundle_count = output_bundle2.get_child_bundle_count()
            for i in range(child_bundle_count):
                child_bundle = output_bundle2.get_child_bundle(i)
                attr = child_bundle.get_attribute_by_name("sourcePrimPath")
                child_bundle_paths.add(attr.get())

            self.assertEqual(child_bundle_paths, cube_paths.difference(expected_prim_paths))

    # ----------------------------------------------------------------------
    async def test_readtime(self):
        """Test omni.graph.nodes.ReadTime"""
        app = omni.kit.app.get_app()
        timeline = omni.timeline.get_timeline_interface()

        time_node_name = "Time"

        # setup initial timeline state and avoid possibility that running this test twice will give different results
        timeline.set_fast_mode(True)
        fps = 24.0
        timeline.set_time_codes_per_second(fps)
        timeline.set_start_time(-1.0)
        timeline.set_end_time(0.0)
        timeline.set_current_time(0.0)
        await app.next_update_async()

        keys = og.Controller.Keys
        (_, (time_node,), _, _) = og.Controller.edit(
            self.TEST_GRAPH_PATH, {keys.CREATE_NODES: [(time_node_name, "omni.graph.nodes.ReadTime")]}
        )
        await app.next_update_async()
        await og.Controller.evaluate()
        self.assertGreater(og.Controller(og.Controller.attribute("outputs:deltaSeconds", time_node)).get(), 0.0)
        self.assertGreater(og.Controller(og.Controller.attribute("outputs:timeSinceStart", time_node)).get(), 0.0)
        self.assertFalse(og.Controller(og.Controller.attribute("outputs:isPlaying", time_node)).get())
        self.assertEqual(og.Controller(og.Controller.attribute("outputs:time", time_node)).get(), 0.0)
        self.assertEqual(og.Controller(og.Controller.attribute("outputs:frame", time_node)).get(), 0.0)

        timeline.set_start_time(1.0)
        timeline.set_end_time(10.0)
        timeline.play()
        await app.next_update_async()
        await app.next_update_async()
        await og.Controller.evaluate()

        # graph has been reloaded as a result of the fps change: we need re refetch the node
        time_node = og.get_node_by_path(self.TEST_GRAPH_PATH + "/" + time_node_name)

        self.assertTrue(og.Controller(og.Controller.attribute("outputs:isPlaying", time_node)).get())
        animation_time = og.Controller(og.Controller.attribute("outputs:time", time_node)).get()
        self.assertGreater(animation_time, 0.0)
        self.assertAlmostEqual(
            og.Controller(og.Controller.attribute("outputs:frame", time_node)).get(), fps * animation_time, places=2
        )
        timeline.stop()
        await app.next_update_async()

    # ----------------------------------------------------------------------
    async def test_read_and_write_prim_material(self):
        await self._read_and_write_prim_material(False)

    async def test_read_and_write_prim_material_with_target(self):
        await self._read_and_write_prim_material(True)

    async def _read_and_write_prim_material(self, use_target_inputs):
        """
        Test ReadPrimMaterial and WritePrimMaterial node. If use_target_inputs is true, use the prim target input rather
        than the prim path input
        """
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        mat_path = "/TestMaterial"
        prim_path_a = "/TestPrimA"
        prim_path_b = "/TestPrimB"

        (graph, (read_node, _), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimMaterial"),
                    ("Write", "omni.graph.nodes.WritePrimMaterial"),
                ],
                keys.CREATE_PRIMS: [(mat_path, "Material"), (prim_path_a, "Cube"), (prim_path_b, "Cube")],
                keys.SET_VALUES: [
                    ("Read.inputs:primPath", prim_path_a),
                    ("Write.inputs:primPath", prim_path_b),
                    ("Write.inputs:materialPath", mat_path),
                ],
            },
        )
        await controller.evaluate(graph)

        if use_target_inputs:
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetRelationshipAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prim"),
                target=prim_path_a,
            )
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetRelationshipAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prim"),
                target=prim_path_b,
            )
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetRelationshipAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:material"),
                target=mat_path,
            )
            await controller.evaluate(graph)

        prim_a = stage.GetPrimAtPath(prim_path_a)
        rel_a = prim_a.CreateRelationship("material:binding")
        omni.kit.commands.execute("AddRelationshipTarget", relationship=rel_a, target=Sdf.Path(mat_path))
        await controller.evaluate(graph)

        read_output = og.Controller.get(controller.attribute("outputs:material", read_node))
        self.assertEqual(read_output, mat_path)

        prim_b = stage.GetPrimAtPath(prim_path_b)
        mat_b, _ = UsdShade.MaterialBindingAPI(prim_b).ComputeBoundMaterial()
        write_output = mat_b.GetPath().pathString
        self.assertEqual(write_output, mat_path)

    # ----------------------------------------------------------------------
    async def test_read_prim_material_with_target_connections(self):
        """Test ReadPrimMaterial and WritePrimMaterial node"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        mat_path = "/TestMaterial"
        prim_path_a = "/TestPrimA"

        (graph, (_, _, prim_node, _), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimMaterial"),
                    ("GetPrimAtPath", "omni.graph.nodes.GetPrimsAtPath"),
                    ("GetPrimPath", "omni.graph.nodes.GetPrimPath"),
                    ("ConstToken", "omni.graph.nodes.ConstantToken"),
                ],
                keys.CREATE_PRIMS: [(mat_path, "Material"), (prim_path_a, "Cube")],
                keys.SET_VALUES: [("ConstToken.inputs:value", prim_path_a)],
                keys.CONNECT: [
                    ("Read.outputs:materialPrim", "GetPrimPath.inputs:prim"),
                    ("GetPrimAtPath.outputs:prims", "Read.inputs:prim"),
                    ("ConstToken.inputs:value", "GetPrimAtPath.inputs:path"),
                ],
            },
        )

        prim_a = stage.GetPrimAtPath(prim_path_a)
        rel_a = prim_a.CreateRelationship("material:binding")
        omni.kit.commands.execute("AddRelationshipTarget", relationship=rel_a, target=Sdf.Path(mat_path))
        await controller.evaluate(graph)

        mat_output = og.Controller.get(("outputs:primPath", prim_node))
        self.assertEqual(mat_output, mat_path)

    # ----------------------------------------------------------------------
    async def test_read_and_write_settings(self):
        """Test ReadSetting and WriteSetting node"""
        controller = og.Controller()
        keys = og.Controller.Keys
        settings = carb.settings.get_settings()

        setting_path = self.PERSISTENT_SETTINGS_PREFIX + "/omnigraph/deprecationsAreErrors"
        default_value = settings.get(setting_path)
        sample_value = not default_value

        (graph, (read_node, _, _), _, _) = og.Controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadSetting"),
                    ("Write", "omni.graph.nodes.WriteSetting"),
                    ("ConstBool", "omni.graph.nodes.ConstantBool"),
                ],
                keys.CONNECT: [("ConstBool.inputs:value", "Write.inputs:value")],
                keys.SET_VALUES: [
                    ("Read.inputs:settingPath", setting_path),
                    ("Write.inputs:settingPath", setting_path),
                    ("ConstBool.inputs:value", sample_value),
                ],
            },
        )
        await controller.evaluate(graph)
        await omni.kit.app.get_app().next_update_async()

        output_value = og.Controller.get(controller.attribute("outputs:value", read_node))
        self.assertEqual(output_value, sample_value)
        settings.set(setting_path, default_value)

    # ----------------------------------------------------------------------
    async def test_get_look_at_rotation(self):
        """Test GetLookAtRotation node"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        # Get the scene up-vector.
        stage_up_axis = stage.GetMetadata("upAxis")
        if stage_up_axis.lower() == "x":
            up_vec = Gf.Vec3d.XAxis()
            fwd_vec = -Gf.Vec3d.ZAxis()
            right_vec = -Gf.Vec3d.YAxis()
            is_y_up = False
        elif stage_up_axis.lower() == "z":
            up_vec = Gf.Vec3d.ZAxis()
            fwd_vec = -Gf.Vec3d.XAxis()
            right_vec = Gf.Vec3d.YAxis()
            is_y_up = False
        else:
            up_vec = Gf.Vec3d.YAxis()
            fwd_vec = -Gf.Vec3d.ZAxis()
            right_vec = Gf.Vec3d.XAxis()
            is_y_up = True

        # Create a prim driven by a GetLookAtRotation node to always face its fwd_vec toward
        # the origin.

        distance_from_origin = 750.0

        # It's important that the up component be non-zero as that creates a difference between the local and
        # and world up-vectors, which is where problems can occur.
        up_offset = 100.0
        position = up_vec * up_offset - fwd_vec * distance_from_origin
        xform_path = "/World/xform"

        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (lookat_node, *_), (xform_prim, *_), _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("lookat", "omni.graph.nodes.GetLookAtRotation"),
                    ("write", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: [(xform_path, "Xform")],
                keys.CONNECT: [("lookat.outputs:rotateXYZ", "write.inputs:value")],
                keys.SET_VALUES: [
                    ("lookat.inputs:forward", fwd_vec),
                    ("lookat.inputs:start", position),
                    ("write.inputs:name", "xformOp:rotateXYZ"),
                    ("write.inputs:primPath", xform_path),
                    ("write.inputs:usePath", True),
                ],
            },
        )

        await controller.evaluate(graph)

        start_attr = lookat_node.get_attribute("inputs:start")
        fwd_attr = lookat_node.get_attribute("inputs:forward")
        up_attr = lookat_node.get_attribute("inputs:up")

        for _ in range(2):
            # Move the transform in a circle around the origin, keeping the up component constant.
            for angle in range(0, 360, 45):
                rad = math.radians(angle)
                position = right_vec * math.sin(rad) * distance_from_origin
                position -= fwd_vec * math.cos(rad) * distance_from_origin
                position += up_vec * up_offset
                start_attr.set(position)

                await controller.evaluate(graph)

                # Get the normalized vector from the current position to the origin.
                pos_vec = -Gf.Vec3d(position).GetNormalized()

                # The forward vector in world space should match the position vector.
                cache = UsdGeom.XformCache()
                xform = cache.GetLocalToWorldTransform(xform_prim)
                ws_fwd = xform.TransformDir(fwd_vec).GetNormalized()
                diff = Gf.Rotation(pos_vec, ws_fwd).GetAngle()
                self.assertAlmostEqual(
                    diff, 0.0, msg=f"at {angle} degree position, forward vector does not point at origin."
                )

                # The local and world up-vectors should remain within 8 degrees of each other.
                ws_up = xform.TransformDir(up_vec).GetNormalized()
                diff = Gf.Rotation(up_vec, ws_up).GetAngle()
                self.assertLess(abs(diff), 8.0, f"at {angle} degree position, up-vector has drifted too far")

            # The first pass used scene-up.
            # For the second pass we set an explicit up-vector, different from scene-up.
            if is_y_up:
                up_vec = Gf.Vec3d.ZAxis()
                fwd_vec = -Gf.Vec3d.XAxis()
                right_vec = Gf.Vec3d.YAxis()
            else:
                up_vec = Gf.Vec3d.YAxis()
                fwd_vec = -Gf.Vec3d.ZAxis()
                right_vec = Gf.Vec3d.XAxis()

            fwd_attr.set(fwd_vec)
            up_attr.set(up_vec)

    async def test_selection_nodes(self):
        """Test ReadStageSelection, IsPrimSelected nodes"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (read_node, check_node, check_target), _, _) = og.Controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadStageSelection"),
                    ("IsSelected", "omni.graph.nodes.IsPrimSelected"),
                    ("IsSelectedWithTarget", "omni.graph.nodes.IsPrimSelected"),
                ],
                keys.SET_VALUES: [("IsSelected.inputs:primPath", f"{self.TEST_GRAPH_PATH}/Read")],
            },
        )
        await controller.evaluate(graph)

        rel = stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/IsSelectedWithTarget.inputs:prim")
        omni.kit.commands.execute("AddRelationshipTarget", relationship=rel, target=f"{self.TEST_GRAPH_PATH}/Read")

        selection = usd_context.get_selection()

        self.assertFalse(og.Controller.get(controller.attribute("outputs:isSelected", check_node)))
        self.assertFalse(og.Controller.get(controller.attribute("outputs:isSelected", check_target)))
        selection.set_selected_prim_paths(
            [read_node.get_prim_path()],
            False,
        )
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()

        selected_paths = og.Controller.get(controller.attribute("outputs:selectedPrims", read_node))
        self.assertListEqual(selected_paths, selection.get_selected_prim_paths())
        self.assertTrue(og.Controller.get(controller.attribute("outputs:isSelected", check_node)))
        self.assertTrue(og.Controller.get(controller.attribute("outputs:isSelected", check_target)))

        selection.clear_selected_prim_paths()
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        selected_paths = og.Controller.get(controller.attribute("outputs:selectedPrims", read_node))
        self.assertListEqual(selected_paths, selection.get_selected_prim_paths())
        self.assertFalse(og.Controller.get(controller.attribute("outputs:isSelected", check_node)))
        self.assertFalse(og.Controller.get(controller.attribute("outputs:isSelected", check_target)))

    # ----------------------------------------------------------------------
    async def test_parallel_read_write_prim_attribute(self):
        """Attempt to read and write prim attributes in parallel"""

        # ReadIndex -> Multiply
        #           /           \
        #         10             ------x--------> MakeTranslate -> WriteTranslate
        #                                      /
        # ReadTranslate -> BreakTranslate -yz-

        keys = og.Controller.Keys

        index_name = "testIndex"
        attr_name = "xformOp:translate"
        scale = 10

        cubes = []

        for i in range(64):
            graph_name = f"/graph{i}"
            obj_name = f"/obj{i}"
            controller = og.Controller()
            (graph, _, prims, _) = controller.edit(
                graph_name,
                {
                    keys.CREATE_NODES: [
                        ("ReadIndex", "omni.graph.nodes.ReadPrimAttribute"),
                        ("ReadTranslate", "omni.graph.nodes.ReadPrimAttribute"),
                        ("BreakTranslate", "omni.graph.nodes.BreakVector3"),
                        ("Multiply", "omni.graph.nodes.Multiply"),
                        ("MakeTranslate", "omni.graph.nodes.MakeVector3"),
                        ("WriteTranslate", "omni.graph.nodes.WritePrimAttribute"),
                    ],
                    keys.CREATE_PRIMS: [
                        (obj_name, "Cube"),
                    ],
                },
            )

            # we want to test parallel execution: prevent any merging
            graph.set_auto_instancing_allowed(False)

            attr = prims[0].CreateAttribute(index_name, Sdf.ValueTypeNames.Double)
            attr.Set(i)
            cubes.append(prims[0])

            controller.edit(
                graph_name,
                {
                    keys.SET_VALUES: [
                        ("ReadIndex.inputs:name", index_name),
                        ("ReadIndex.inputs:usePath", True),
                        ("ReadIndex.inputs:primPath", obj_name),
                        ("ReadTranslate.inputs:name", attr_name),
                        ("ReadTranslate.inputs:usePath", True),
                        ("ReadTranslate.inputs:primPath", obj_name),
                        ("WriteTranslate.inputs:name", attr_name),
                        ("WriteTranslate.inputs:usePath", True),
                        ("WriteTranslate.inputs:primPath", obj_name),
                        ("Multiply.inputs:b", scale, "double"),
                    ],
                },
            )

            controller.edit(
                graph_name,
                {
                    keys.CONNECT: [
                        ("ReadIndex.outputs:value", "Multiply.inputs:a"),
                        ("ReadTranslate.outputs:value", "BreakTranslate.inputs:tuple"),
                        ("Multiply.outputs:product", "MakeTranslate.inputs:x"),
                        ("BreakTranslate.outputs:y", "MakeTranslate.inputs:y"),
                        ("BreakTranslate.outputs:z", "MakeTranslate.inputs:z"),
                        ("MakeTranslate.outputs:tuple", "WriteTranslate.inputs:value"),
                    ],
                },
            )

        # cubes[0].GetStage().Export("c:/tmp/test.usda")

        for cube in cubes:
            translate = cube.GetAttribute(attr_name).Get()
            self.assertEqual(translate[0], 0)

        await og.Controller.evaluate()

        for cube in cubes:
            index = cube.GetAttribute(index_name).Get()
            translate = cube.GetAttribute(attr_name).Get()
            self.assertEqual(translate[0], index * scale)

    # ----------------------------------------------------------------------
    async def test_read_prim_node_default_time_code_upgrades(self):
        """
        Tests that read prim node variants property upgrade their time code property
        to use NAN as a default
        """
        app = omni.kit.app.get_app()

        nodes = [
            "ReadPrim",
            "ReadPrims",
            "ReadPrimAttribute",
            "ReadPrimAttributes",
            "ReadPrimBundle",
            "ReadPrimsBundle",
        ]

        # testing under different frame rates to validate the upgrade is correctly detected
        frame_rates = [12, 24, 30, 60]
        for frame_rate in frame_rates:
            with self.subTest(frame_rate=frame_rate):
                timeline = omni.timeline.get_timeline_interface()
                timeline.set_time_codes_per_second(frame_rate)
                await app.next_update_async()

                file_path = os.path.join(os.path.dirname(__file__), "data", "TestReadPrimNodeVariants.usda")
                stage = omni.usd.get_context().get_stage()

                # create a prim and add a file reference
                prim = stage.DefinePrim("/Ref")
                # when we make the references, there will be deprecation warnings
                prim.GetReferences().AddReference(file_path)

                self.assertEqual(stage.GetTimeCodesPerSecond(), frame_rate)

                await omni.kit.app.get_app().next_update_async()

                for node in nodes:
                    attr = og.Controller.attribute(f"/Ref/ActionGraph/{node}.inputs:usdTimecode")
                    self.assertTrue(isnan(attr.get()))

                await omni.usd.get_context().new_stage_async()

    # ----------------------------------------------------------------------
    async def test_is_prim_active(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        contents = Usd.Stage.CreateInMemory("payload.usd")
        contents.DefinePrim("/payload/A", "Xform")
        payload = stage.DefinePrim("/payload", "Xform")
        payload.GetPayloads().AddPayload(Sdf.Payload(contents.GetRootLayer().identifier, "/payload"))

        keys = og.Controller.Keys
        controller = og.Controller()
        (graph, (node,), _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [("IsActive", "omni.graph.nodes.IsPrimActive")],
            },
        )
        await controller.evaluate(graph)

        rel = stage.GetRelationshipAtPath("/TestGraph/IsActive.inputs:primTarget")
        omni.kit.commands.execute("AddRelationshipTarget", relationship=rel, target="/payload/A")
        await controller.evaluate(graph)

        out = node.get_attribute("outputs:active")
        self.assertEqual(out.get(), True)

        # Note: currently the node will raise an error if the prim isn't in the stage as well as set active to false
        payload.Unload()
        with ogts.ExpectedError():
            await controller.evaluate(graph)
        out = node.get_attribute("outputs:active")
        self.assertEqual(out.get(), False)

    # -------------------------------------------------------------------------------
    async def test_get_graph_target_prim(self):
        """Test that the GetGraphTargetPrim node outputs the expected graph target prim"""

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_graph_target_node, path_node), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("GetGraphTargetPrim", "omni.graph.nodes.GetGraphTargetPrim"),
                    ("GetPrimPath", "omni.graph.nodes.GetPrimPath"),
                ],
                keys.CONNECT: [("GetGraphTargetPrim.outputs:prim", "GetPrimPath.inputs:prim")],
            },
        )

        # no instances, it should return the graph prim
        await og.Controller.evaluate(graph)
        self.assertEqual(self.TEST_GRAPH_PATH, path_node.get_attribute("outputs:primPath").get())

        num_instances = 5

        # create instances
        for i in range(0, num_instances):
            prim_name = f"/World/Prim_{i}"
            stage.DefinePrim(prim_name)
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim_name, self.TEST_GRAPH_PATH)

        # rerun the graph, this time evaluating instances
        await og.Controller.evaluate(graph)

        for i in range(0, num_instances):
            self.assertEqual(f"/World/Prim_{i}", path_node.get_attribute("outputs:primPath").get(instance=i))

    async def test_construct_array_invalid_array_size(self):
        """Test that setting an invalid value on OgnConstructArray.arraySize does not crash"""
        # Ideally this would be done with the following test in the .ogn file, but the node parser
        # applies the minimum / maximum attribute limits when parsing the tests, which means that
        # we cannot define tests for invalid inputs in the .ogn file
        # {
        #     "inputs:arraySize": -1, "inputs:arrayType": "Int64",
        #     "outputs:array": {"type": "int64[]", "value": []}
        # }
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, nodes, _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [("ConstructArray", "omni.graph.nodes.ConstructArray")],
                keys.SET_VALUES: [
                    ("ConstructArray.inputs:arraySize", -1),
                    ("ConstructArray.inputs:arrayType", "int[]"),
                ],
            },
        )

        # Evaluate the graph once to ensure the node has a Database created
        await controller.evaluate(graph)

        result = nodes[0].get_attribute("outputs:array").get()
        expected: [int] = []
        self.assertListEqual(expected, list(result))

    async def test_get_prim_local_to_world_empty(self):
        """Regression test for GetPrimLocalToWorldTransform"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [("TestWorld", "omni.graph.nodes.GetPrimLocalToWorldTransform")],
                keys.SET_VALUES: [
                    ("TestWorld.inputs:usePath", False),
                ],
            },
        )
        with ogts.ExpectedError():
            await controller.evaluate(graph)
