# noqa: PLC0302
"""Tests for the various OmniGraph commands"""
from typing import List, Tuple

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.commands
import omni.kit.test
import omni.usd
from pxr import OmniGraphSchemaTools, Usd


class TestTopologyCommands(ogts.OmniGraphTestCase):

    # ----------------------------------------------------------------------

    async def test_create_graph_global(self):
        """
        test the CreateGraphAsNodeCommand with global graphs
        """
        # This is the main command we'll use
        # to create new graphs for users going forward.  These graphs are top level graphs, and are wrapped by nodes
        # in an overall global container graph
        # We also leverage this setup to test the creation of dynamic attributes with USD backing

        orchestration_graphs = og.get_global_orchestration_graphs()
        # 4, because we have 3 separate pipeline stages currently (simulation, pre-render, post-render, on-demand)
        self.assertEqual(len(orchestration_graphs), 4)

        orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(
            og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_POSTRENDER
        )
        self.assertEqual(len(orchestration_graphs), 1)
        self.assertEqual(orchestration_graphs[0].get_path_to_graph(), "/__postRenderRootGraph")
        self.assertEqual(
            orchestration_graphs[0].get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_POSTRENDER
        )
        self.assertEqual(orchestration_graphs[0].get_graph_backing_type(), og.GraphBackingType.GRAPH_BACKING_TYPE_NONE)
        self.assertEqual(
            orchestration_graphs[0].evaluation_mode, og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_AUTOMATIC
        )

        orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(
            og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION
        )
        self.assertEqual(len(orchestration_graphs), 1)
        orchestration_graph = orchestration_graphs[0]
        (result, wrapper_node) = og.cmds.CreateGraphAsNode(
            graph=orchestration_graph,
            node_name="my_push_graph",
            graph_path="/World/my_push_graph",
            evaluator_name="push",
            is_global_graph=True,
            backed_by_usd=True,
            fc_backing_type=og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED,
            pipeline_stage=og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION,
            evaluation_mode=og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_STANDALONE,
        )
        self.assertTrue(result)
        self.assertTrue(wrapper_node.is_valid())
        self.assertEqual(wrapper_node.get_prim_path(), "/__simulationRootGraph/_World_my_push_graph")
        wrapped_graph = wrapper_node.get_wrapped_graph()
        self.assertTrue(wrapped_graph.is_valid())
        self.assertEqual(wrapped_graph.get_path_to_graph(), "/World/my_push_graph")

        self.assertEqual(wrapped_graph.get_graph_backing_type(), og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED)
        self.assertEqual(wrapped_graph.get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION)
        self.assertEqual(wrapped_graph.evaluation_mode, og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_STANDALONE)

        (_, (_, counter_node), _, _) = og.Controller.edit(
            wrapped_graph,
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                og.Controller.Keys.SET_VALUES: [("OnTick.inputs:onlyPlayback", False)],
                og.Controller.Keys.CONNECT: [
                    ("OnTick.outputs:tick", "Counter.inputs:execIn"),
                ],
            },
        )

        # What we should see in response is that the counter is incrementing
        await og.Controller.evaluate(wrapped_graph)
        outputs_attr_value1 = og.Controller.get(counter_node.get_attribute("outputs:count"))

        await og.Controller.evaluate(wrapped_graph)
        outputs_attr_value2 = og.Controller.get(counter_node.get_attribute("outputs:count"))

        self.assertGreater(outputs_attr_value2, outputs_attr_value1)

        # test dynamic attribute creation, with USD backing:
        # try setting a simple value and retrieving it
        counter_node.create_attribute(
            "inputs:test_create1",
            og.Type(og.BaseDataType.INT),
            og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
            None,
            og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_REGULAR,
            "",
        )
        test_create1_attr = counter_node.get_attribute("inputs:test_create1")
        self.assertTrue(test_create1_attr.is_valid())
        self.assertTrue(test_create1_attr.is_dynamic())
        og.cmds.SetAttr(attr=test_create1_attr, value=123)
        test_create1_attr_value = og.Controller.get(test_create1_attr)
        self.assertEqual(test_create1_attr_value, 123)
        self.assertEqual(test_create1_attr.get_port_type(), og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT)

        # try setting a tuple value and retrieving it
        counter_node.create_attribute(
            "inputs:test_create2",
            og.Type(og.BaseDataType.INT, 3),
            og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
            None,
            og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_REGULAR,
            "",
        )
        test_create2_attr = counter_node.get_attribute("inputs:test_create2")
        self.assertTrue(test_create2_attr.is_valid())
        og.cmds.SetAttr(attr=test_create2_attr, value=(1, 2, 3))
        test_create2_attr_value = og.Controller.get(test_create2_attr)
        self.assertEqual(test_create2_attr_value[0], 1)
        self.assertEqual(test_create2_attr_value[1], 2)
        self.assertEqual(test_create2_attr_value[2], 3)

        # try setting a tuple array value and retrieving it
        counter_node.create_attribute(
            "inputs:test_create3",
            og.Type(og.BaseDataType.INT, 3, 1),
            og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
            None,
            og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_REGULAR,
            "",
        )
        test_create3_attr = counter_node.get_attribute("inputs:test_create3")
        self.assertTrue(test_create3_attr.is_valid())
        og.cmds.SetAttr(attr=test_create3_attr, value=[[1, 2, 3], [4, 5, 6]])
        test_create3_attr_value = og.Controller.get(test_create3_attr)
        v1 = test_create3_attr_value[1]
        self.assertEqual(v1[0], 4)
        self.assertEqual(v1[1], 5)
        self.assertEqual(v1[2], 6)

        # try creating an execution attribute
        counter_node.create_attribute(
            "inputs:test_create4",
            og.Type(og.BaseDataType.UINT, 1, 0, og.AttributeRole.EXECUTION),
            og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
            None,
            og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_REGULAR,
            "",
        )
        test_create4_attr = counter_node.get_attribute("inputs:test_create4")
        self.assertTrue(test_create4_attr.is_valid())

    # ----------------------------------------------------------------------
    async def test_undo_redo_connection(self):
        """Test undo and redo of node creation and connection"""

        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, _, _, _) = controller.edit("/TestGraph")

        async def create_and_connect(ix):
            controller.edit(
                graph,
                {
                    keys.CREATE_NODES: [
                        (f"SimpleA{ix}", "omni.graph.tutorials.SimpleData"),
                        (f"SimpleB{ix}", "omni.graph.tutorials.SimpleData"),
                    ]
                },
            )
            await controller.evaluate(graph)

            controller.edit(graph, {keys.CONNECT: (f"SimpleA{ix}.outputs:a_float", f"SimpleB{ix}.inputs:a_float")})
            await controller.evaluate(graph)

        await create_and_connect(1)
        # undo the connection
        omni.kit.undo.undo()
        # undo the node creation
        omni.kit.undo.undo()
        # redo node creation
        omni.kit.undo.redo()
        # redo connection
        omni.kit.undo.redo()
        # undo the connection
        omni.kit.undo.undo()
        # undo the node creation
        omni.kit.undo.undo()

    # --------------------------------------------------------------------------------------------------------------
    def __confirm_connection_counts(self, attribute: og.Attribute, upstream_count: int, downstream_count: int):
        """Check that the given attribute has the proscribed number of upstream and downstream connections"""
        self.assertEqual(attribute.get_upstream_connection_count(), upstream_count)
        self.assertEqual(attribute.get_downstream_connection_count(), downstream_count)

    # ----------------------------------------------------------------------
    async def test_disconnect_all_attrs(self):
        """
        Test the DisconnectAllAttrsCommand. This command takes a single attribute as parameter and disconnects all
        incoming and outgoing connections on it.
        """
        controller = og.Controller()
        keys = og.Controller.Keys
        #
        # Source --> Sink    --> InOut
        #       \    \      /
        #        \    -----/
        #         --> FanOut
        (graph, (source_node, sink_node, fanout_node, inout_node), _, _,) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("Source", "omni.graph.tutorials.SimpleData"),
                    ("Sink", "omni.graph.tutorials.SimpleData"),
                    ("FanOut", "omni.graph.tutorials.SimpleData"),
                    ("InOut", "omni.graph.tutorials.SimpleData"),
                ],
                keys.CONNECT: [
                    ("Source.outputs:a_bool", "Sink.inputs:a_bool"),
                    ("Source.outputs:a_bool", "FanOut.inputs:a_bool"),
                    ("Sink.inputs:a_bool", "InOut.inputs:a_bool"),
                ],
            },
        )

        def _get_attributes():
            return (
                og.Controller.attribute("outputs:a_bool", ("Source", graph)),
                og.Controller.attribute("inputs:a_bool", ("Sink", graph)),
                og.Controller.attribute("inputs:a_bool", ("FanOut", graph)),
                og.Controller.attribute("inputs:a_bool", ("InOut", graph)),
            )

        source_output, sink_input, fanout_input, inout_input = _get_attributes()

        # The sequence tests successful restoration when everything is deleted and restored in groups;
        #    DisconnectAll(Source.outputs:a_bool)
        #        Removes Source->Sink and Source->FanOut
        #    DisconnectAll(Sink.inputs:a_bool)
        #        Removes Sink->InOut
        #    Delete Source, Sink, InOut, FanOut
        #    undo deletes
        #    undo second disconnect
        #        Restores Sink->InOut
        #    undo first disconnect
        #        Restores Source->Sink and Source->FanOut
        #    redo and undo to go the beginning and get back here again
        #    DisconnectAll(Sink.inputs:a_bool)
        #        Removes Sink->InOut and Source->Sink
        #    undo

        og.cmds.DisconnectAllAttrs(attr=source_output, modify_usd=True)
        self.__confirm_connection_counts(source_output, 0, 0)
        self.__confirm_connection_counts(sink_input, 0, 1)
        self.__confirm_connection_counts(fanout_input, 0, 0)
        self.__confirm_connection_counts(inout_input, 1, 0)

        og.cmds.DisconnectAllAttrs(attr=sink_input, modify_usd=True)
        self.__confirm_connection_counts(source_output, 0, 0)
        self.__confirm_connection_counts(sink_input, 0, 0)
        self.__confirm_connection_counts(fanout_input, 0, 0)
        self.__confirm_connection_counts(inout_input, 0, 0)

        og.cmds.DeleteNode(graph=graph, node_path=source_node.get_prim_path(), modify_usd=True)
        og.cmds.DeleteNode(graph=graph, node_path=sink_node.get_prim_path(), modify_usd=True)
        og.cmds.DeleteNode(graph=graph, node_path=fanout_node.get_prim_path(), modify_usd=True)
        og.cmds.DeleteNode(graph=graph, node_path=inout_node.get_prim_path(), modify_usd=True)

        omni.kit.undo.undo()
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        # May have lost the objects through the undo process so get them again
        source_output, sink_input, fanout_input, inout_input = _get_attributes()
        self.__confirm_connection_counts(source_output, 0, 0)
        self.__confirm_connection_counts(sink_input, 0, 1)
        self.__confirm_connection_counts(fanout_input, 0, 0)
        self.__confirm_connection_counts(inout_input, 1, 0)
        omni.kit.undo.undo()
        self.__confirm_connection_counts(source_output, 0, 2)
        self.__confirm_connection_counts(sink_input, 1, 1)
        self.__confirm_connection_counts(fanout_input, 1, 0)
        self.__confirm_connection_counts(inout_input, 1, 0)

        omni.kit.undo.redo()
        omni.kit.undo.redo()
        omni.kit.undo.redo()
        omni.kit.undo.redo()
        omni.kit.undo.redo()
        omni.kit.undo.redo()
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        # May have lost the objects through the undo process so get them again
        source_output, sink_input, fanout_input, inout_input = _get_attributes()
        self.__confirm_connection_counts(source_output, 0, 2)
        self.__confirm_connection_counts(sink_input, 1, 1)
        self.__confirm_connection_counts(fanout_input, 1, 0)
        self.__confirm_connection_counts(inout_input, 1, 0)

        og.cmds.DisconnectAllAttrs(attr=sink_input, modify_usd=True)
        self.__confirm_connection_counts(source_output, 0, 1)
        self.__confirm_connection_counts(sink_input, 0, 0)
        self.__confirm_connection_counts(fanout_input, 1, 0)
        self.__confirm_connection_counts(inout_input, 0, 0)

        omni.kit.undo.undo()
        self.__confirm_connection_counts(source_output, 0, 2)
        self.__confirm_connection_counts(sink_input, 1, 1)
        self.__confirm_connection_counts(fanout_input, 1, 0)
        self.__confirm_connection_counts(inout_input, 1, 0)

    # ----------------------------------------------------------------------
    async def test_change_pipeline_stage_invalid(self):
        """
        Test the ChangePipelineStageCommand when given an invalid pipeline stage
        """
        controller = og.Controller()
        (graph, _, _, _) = controller.edit("/World/TestGraph")
        self.assertEqual(graph.get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION)

        # Check that we are not allowed to set pipeline stage to unknown
        og.cmds.ChangePipelineStage(graph=graph, new_pipeline_stage=og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_UNKNOWN)
        self.assertEqual(graph.get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION)

        # Check that setting the same pipeline stage does nothing
        og.cmds.ChangePipelineStage(
            graph=graph, new_pipeline_stage=og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION
        )
        self.assertEqual(graph.get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION)

    # -------------------------------------------------------------------------
    async def next_frame(self, node) -> bool:
        """Helper method that ticks the next frame, returning how many times the node was ticked"""
        start_count = node.get_compute_count()
        await omni.kit.app.get_app().next_update_async()
        end_count = node.get_compute_count()
        return end_count - start_count

    # -------------------------------------------------------------------------
    async def test_evaluation_modes_command(self):
        """Test evaluation modes using instanced and non-instanced graphs"""

        _graph_path = "/World/TestGraph"
        __graph_count = 0

        def create_simple_graph(path) -> Tuple[og.Graph, List[og.Node]]:
            nonlocal __graph_count
            controller = og.Controller()
            keys = og.Controller.Keys
            (graph, nodes, _, _) = controller.edit(
                f"{path}_{__graph_count}",
                {
                    keys.CREATE_NODES: [
                        ("OnTick", "omni.graph.action.OnTick"),
                    ]
                },
            )
            __graph_count += 1
            return (graph, nodes)

        def create_graph_instance(name: str, graph: og.Graph) -> Usd.Prim:
            """Creates an instance of the given graph"""
            stage = omni.usd.get_context().get_stage()
            prim = stage.DefinePrim(name)
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim.GetPath(), graph.get_path_to_graph())
            return prim

        stage = omni.usd.get_context().get_stage()
        (graph, nodes) = create_simple_graph(_graph_path)

        timeline = omni.timeline.get_timeline_interface()
        timeline.set_target_framerate(timeline.get_time_codes_per_seconds())
        timeline.set_start_time(0)
        timeline.set_end_time(1000)
        timeline.play()

        node = nodes[0]
        self.assertEquals(graph.evaluation_mode, og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_AUTOMATIC)

        # in standalone mode, the node should tick
        ticks = await self.next_frame(node)
        self.assertEquals(ticks, 1)

        # in asset mode it should not
        og.cmds.SetEvaluationMode(
            graph=graph, new_evaluation_mode=og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_INSTANCED
        )
        ticks = await self.next_frame(node)
        self.assertEquals(ticks, 0)

        # create graph instances
        prims = [create_graph_instance(f"/World/Prim{i}", graph) for i in range(0, 2)]

        # automatic mode reference mode should evaluate once (vectorized) for all instances
        # the returned compute count should be equal to the number of instances though
        expected_instanced_ticks = 2

        ticks = await self.next_frame(node)
        self.assertEquals(ticks, expected_instanced_ticks)

        og.cmds.SetEvaluationMode(
            graph=graph, new_evaluation_mode=og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_AUTOMATIC
        )
        ticks = await self.next_frame(node)
        self.assertEquals(ticks, expected_instanced_ticks)

        # standalone mode should run itself and not the prims
        og.cmds.SetEvaluationMode(
            graph=graph, new_evaluation_mode=og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_STANDALONE
        )
        ticks = await self.next_frame(node)
        self.assertEquals(ticks, 1)

        # delete the prims, standalone should now run
        for p in prims:
            stage.RemovePrim(p.GetPath())

        ticks = await self.next_frame(node)
        self.assertEquals(ticks, 1)

        omni.timeline.get_timeline_interface().stop()

    # ----------------------------------------------------------------------
    async def test_evaluation_modes_command_undo_redo(self):
        """Tests the evaluation mode command has the correct value when undo and redo are applied"""
        controller = og.Controller()
        (graph, _, _, _) = controller.edit("/TestGraph")

        # validate the correct default
        default_value = og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_AUTOMATIC
        self.assertEquals(graph.evaluation_mode, og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_AUTOMATIC)

        for value in [
            og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_AUTOMATIC,
            og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_STANDALONE,
            og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_INSTANCED,
        ]:
            og.cmds.SetEvaluationMode(graph=graph, new_evaluation_mode=value)
            self.assertEquals(graph.evaluation_mode, value)
            omni.kit.undo.undo()
            self.assertEquals(graph.evaluation_mode, default_value)
            omni.kit.undo.redo()
            self.assertEquals(graph.evaluation_mode, value)
            omni.kit.undo.undo()
            self.assertEquals(graph.evaluation_mode, default_value)

    # ----------------------------------------------------------------------
    async def test_variable_commands(self):
        controller = og.Controller()
        (graph, _, _, _) = controller.edit("/TestGraph")

        variable_name = "TestVariable"
        og.cmds.CreateVariable(
            graph=graph,
            variable_name=variable_name,
            variable_type=og.Type(og.BaseDataType.INT),
            graph_context=graph.get_default_graph_context(),
            variable_value=5,
        )
        variable = graph.find_variable(variable_name)
        self.assertTrue(variable is not None)
        self.assertTrue(variable.valid)
        self.assertEqual(variable.type, og.Type(og.BaseDataType.INT))
        self.assertEqual(og.Controller.get_variable_default_value(variable), 5)

        # undo the command and verify that the variable is no longer in the graph
        omni.kit.undo.undo()
        variable = graph.find_variable(variable_name)
        self.assertTrue(variable is None)

        # redo the command and verify that the variable has been restored in the graph
        omni.kit.undo.redo()
        variable = graph.find_variable(variable_name)
        self.assertTrue(variable is not None)
        self.assertTrue(variable.valid)
        self.assertEqual(variable.type, og.Type(og.BaseDataType.INT))
        self.assertEqual(og.Controller.get_variable_default_value(variable), 5)

        test_tooltip = "test tooltip"
        og.cmds.SetVariableTooltip(variable=variable, tooltip=test_tooltip)
        self.assertEqual(variable.tooltip, test_tooltip)

        # undo the command and verify that the variable tooltip is cleared
        omni.kit.undo.undo()
        self.assertEqual(variable.tooltip, "")

        # redo the command and verify that the variable tooltip is restored
        omni.kit.undo.redo()
        self.assertEqual(variable.tooltip, test_tooltip)

        og.cmds.RemoveVariable(graph=graph, variable=variable, graph_context=graph.get_default_graph_context())

        variable = graph.find_variable(variable_name)
        self.assertTrue(variable is None)

        # undo the command and verify that the variable has been restored in the graph
        omni.kit.undo.undo()
        variable = graph.find_variable(variable_name)
        self.assertTrue(variable is not None)
        self.assertTrue(variable.valid)
        self.assertEqual(variable.type, og.Type(og.BaseDataType.INT))
        self.assertEqual(og.Controller.get_variable_default_value(variable), 5)

        # redo the command and verify that the variable is no longer in the graph
        omni.kit.undo.redo()
        variable = graph.find_variable(variable_name)
        self.assertTrue(variable is None)

    # ----------------------------------------------------------------------
    async def test_change_pipeline_stage(self):
        """
        Test the ChangePipelineStageCommand
        """
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, create_nodes, _, _) = controller.edit(
            "/World/TestGraph",
            {
                keys.EXPOSE_PRIMS: [
                    (og.Controller.PrimExposureType.AS_WRITABLE, "/World/Cube", "Write"),
                ],
                keys.CREATE_PRIMS: [
                    ("/World/Cube", {"size": ("int", 1)}),
                ],
                keys.CREATE_NODES: [
                    ("Add", "omni.graph.nodes.Add"),
                    ("Read", "omni.graph.nodes.ReadPrimsBundle"),
                    ("ExtractPrim", "omni.graph.nodes.ExtractPrim"),
                    ("ExtractBundle", "omni.graph.nodes.ExtractBundle"),
                ],
                keys.SET_VALUES: [("ExtractPrim.inputs:primPath", "/World/Cube")],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "ExtractPrim.inputs:prims"),
                    ("ExtractPrim.outputs_primBundle", "ExtractBundle.inputs:bundle"),
                ],
            },
        )

        stage = omni.usd.get_context().get_stage()
        input_prim_prop = stage.GetPropertyAtPath("/World/TestGraph/Read.inputs:prims")
        omni.kit.commands.execute("AddRelationshipTarget", relationship=input_prim_prop, target="/World/Cube")

        await controller.evaluate(graph)

        # check the correctness of the output bundles
        read_node = create_nodes[1]
        extract_prim = create_nodes[2]
        extract_node = create_nodes[3]

        # IBundle2 factory
        factory = og.IBundleFactory.create()

        # Read Prim outputs multiple primitives in bundle
        read_bundle = graph.get_default_graph_context().get_output_bundle(read_node, "outputs_primsBundle")
        self.assertTrue(read_bundle.is_valid())
        read_bundle2 = factory.get_bundle(graph.get_default_graph_context(), read_bundle)
        self.assertTrue(read_bundle2.valid)

        # check number of children from Read Prim
        self.assertEqual(read_bundle2.get_child_bundle_count(), 1)

        # Extract Bundle gets content of the child bundle from Read Prim
        extract_bundle = graph.get_default_graph_context().get_bundle(
            "/World/TestGraph/ExtractBundle/outputs_passThrough"
        )
        self.assertTrue(extract_bundle.is_valid())
        extract_bundle2 = factory.get_bundle(graph.get_default_graph_context(), extract_bundle)
        self.assertTrue(extract_bundle2.valid)

        extract_bundle = graph.get_default_graph_context().get_output_bundle(extract_node, "outputs_passThrough")
        self.assertTrue(extract_bundle.is_valid())
        extract_bundle2 = factory.get_bundle(graph.get_default_graph_context(), extract_bundle)
        self.assertTrue(extract_bundle2.valid)

        # Check number of children and attributes for Extract Bundle
        self.assertEqual(extract_bundle2.get_child_bundle_count(), 0)
        attrib_names = extract_bundle2.get_attribute_names()
        self.assertTrue(len(attrib_names) > 0)

        extract_node_controller = og.Controller(og.Controller.attribute("inputs:primPath", extract_prim))

        controller.edit(
            "/World/TestGraph",
            {
                keys.CONNECT: [
                    ("ExtractBundle.outputs:size", "Add.inputs:a"),
                    ("ExtractBundle.outputs:size", "Add.inputs:b"),
                    ("Add.outputs:sum", "Write.inputs:size"),
                ]
            },
        )

        extract_node = create_nodes[3]
        extract_node_controller = og.Controller(og.Controller.attribute("outputs:size", extract_node))

        # Initially the graph is in simulation pipeline stage
        self.assertEqual(graph.get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION)
        # Check that graph.evaluate() causes graph evaluation
        graph.evaluate()
        self.assertEqual(extract_node_controller.get(), 1)
        graph.evaluate()
        self.assertEqual(extract_node_controller.get(), 2)
        graph.evaluate()
        self.assertEqual(extract_node_controller.get(), 4)
        # Check that omni.kit.app.get_app().next_update_async() also causes graph evaluation
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(extract_node_controller.get(), 8)
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(extract_node_controller.get(), 16)

        # Now change the pipeline stage to on-demand
        og.cmds.ChangePipelineStage(graph=graph, new_pipeline_stage=og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND)
        self.assertEqual(graph.get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND)
        # Check that graph.evaluate() still causes graph evaluation
        graph.evaluate()
        self.assertEqual(extract_node_controller.get(), 32)
        graph.evaluate()
        self.assertEqual(extract_node_controller.get(), 64)
        # Check that omni.kit.app.get_app().next_update_async() does NOT cause graph evaluation
        # because the graph is in on-demand pipeline stage
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(extract_node_controller.get(), 64)
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(extract_node_controller.get(), 64)

        # Undo the command to go back to the simulation pipeline stage
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        omni.kit.undo.undo()
        self.assertEqual(graph.get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION)
        # Check that graph.evaluate() causes graph evaluation
        graph.evaluate()
        self.assertEqual(extract_node_controller.get(), 128)
        # Check that omni.kit.app.get_app().next_update_async() also causes graph evaluation
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(extract_node_controller.get(), 256)

    # ----------------------------------------------------------------------
    async def test_ondemand_pipeline(self):
        """
        test the on-demand pipeline stage.  The on-demand pipeline stage allows graphs to be added to it that tick at
        unknown times instead of the fixed times of graph evaluation of other named stages.  So here we will create
        some graphs in the stage, and manually evaluate it to make sure it works.
        """

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        orchestration_graphs = og.get_global_orchestration_graphs()

        orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(
            og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND
        )
        self.assertEqual(len(orchestration_graphs), 1)
        self.assertEqual(orchestration_graphs[0].get_path_to_graph(), "/__onDemandRootGraph")
        self.assertEqual(
            orchestration_graphs[0].get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND
        )
        self.assertEqual(orchestration_graphs[0].get_graph_backing_type(), og.GraphBackingType.GRAPH_BACKING_TYPE_NONE)

        orchestration_graph = orchestration_graphs[0]
        (result, wrapper_node) = og.cmds.CreateGraphAsNode(
            graph=orchestration_graph,
            node_name="my_push_graph",
            graph_path="/World/my_push_graph",
            evaluator_name="push",
            is_global_graph=True,
            backed_by_usd=True,
            fc_backing_type=og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED,
            pipeline_stage=og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND,
        )
        self.assertTrue(result)
        self.assertTrue(wrapper_node.is_valid())
        self.assertEqual(wrapper_node.get_prim_path(), "/__onDemandRootGraph/_World_my_push_graph")
        wrapped_graph = wrapper_node.get_wrapped_graph()
        self.assertTrue(wrapped_graph.is_valid())
        self.assertEqual(wrapped_graph.get_path_to_graph(), "/World/my_push_graph")

        self.assertEqual(wrapped_graph.get_graph_backing_type(), og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED)
        self.assertEqual(wrapped_graph.get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND)

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

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

        (graph, (_, write_node, _, _, extract_node), _, _) = controller.edit(
            "/World/my_push_graph",
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrims"),
                    ("Write", "omni.graph.nodes.WritePrim"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("ExtractPrim", "omni.graph.nodes.ExtractPrim"),
                    ("ExtractBundle", "omni.graph.nodes.ExtractBundle"),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "ExtractPrim.inputs:prims"),
                    ("ExtractPrim.outputs_primBundle", "ExtractBundle.inputs:bundle"),
                ],
                keys.SET_VALUES: [
                    ("ExtractPrim.inputs:primPath", str(cube_prim.GetPath())),
                ],
            },
        )
        wrapped_graph.evaluate()

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath("/World/my_push_graph/Read.inputs:prims"),
            target=cube_prim.GetPath(),
        )
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath("/World/my_push_graph/Write.inputs:prim"),
            target=cube_prim.GetPath(),
        )
        wrapped_graph.evaluate()

        # node:type
        # node:typeVersion
        # inputs:bundle
        # outputs: passThrough
        n_static_attribs_extract = 4

        # Check we have the expected dynamic attrib on ExtractBundle
        attribs = extract_node.get_attributes()

        bundle = wrapped_graph.get_default_graph_context().get_output_bundle(extract_node, "outputs_passThrough")
        self.assertTrue(bundle.is_valid())

        self.assertEqual(len(attribs) - n_static_attribs_extract, bundle.get_attribute_data_count())

        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)

        # 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(
            graph,
            {
                keys.CONNECT: [
                    ("/World/my_push_graph/ExtractBundle.outputs:size", "/World/my_push_graph/Add.inputs:a"),
                    ("/World/my_push_graph/ExtractBundle.outputs:size", "/World/my_push_graph/Add.inputs:b"),
                    ("/World/my_push_graph/Add.outputs:sum", "/World/my_push_graph/Write.inputs:size"),
                ]
            },
        )
        wrapped_graph.evaluate()
        extract_node_controller = og.Controller(og.Controller.attribute("outputs:size", extract_node))
        self.assertEqual(extract_node_controller.get(), 1)
        wrapped_graph.evaluate()
        self.assertEqual(extract_node_controller.get(), 2)
        wrapped_graph.evaluate()
        self.assertEqual(extract_node_controller.get(), 4)

        # Verify that the graph is not still evaluating while we are not pumping the evaluate
        # OM-4258
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(extract_node_controller.get(), 4)
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(extract_node_controller.get(), 4)

    # ----------------------------------------------------------------------
    async def test_ondemand_pipeline_part_two(self):
        """
        More validation tests for the on-demand pipeline stage. This test was added
        as part of the changes described in OM-87603 to ensure that the pipeline stage
        behaves correctly now that it is executed via the EF code path (as opposed to
        utilizing the legacy schedulers).
        """

        # Define a utility method for creating a simple graph of variable
        # evaluator type and pipeline stage.
        def create_graph(graph_path: str, evaluator_name: str, pipeline: og.GraphPipelineStage):
            """Create a simple graph that increments a counter on a node every time the graph ticks."""
            (graph, nodes, _, _) = og.Controller.edit(
                {"graph_path": graph_path, "evaluator_name": evaluator_name},
                {
                    og.Controller.Keys.CREATE_NODES: [
                        ("OnTick", "omni.graph.action.OnTick"),
                        ("Counter", "omni.graph.action.Counter"),
                    ],
                    og.Controller.Keys.SET_VALUES: [("OnTick.inputs:onlyPlayback", False)],
                    og.Controller.Keys.CONNECT: ("OnTick.outputs:tick", "Counter.inputs:execIn"),
                },
            )

            og.cmds.ChangePipelineStage(graph=graph, new_pipeline_stage=pipeline)
            self.assertEqual(graph.get_pipeline_stage(), pipeline)

            return graph, nodes

        # Currently-supported pipeline stages.
        pipeline_stages = [
            og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION,
            og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_PRERENDER,
            og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_POSTRENDER,
            og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND,
        ]

        # Create a push graph that subscribes to the on-demand pipeline.
        # Doing so should result in the initial population of the OnDemand
        # stage. Check that it only computes when an explicit request to do
        # so is made, and that only the OnDemand orchestration graph has
        # any graph nodes (in this case one).
        _, push_nodes = create_graph("/PushGraph", "push", og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND)
        push_out_cnt_attr = push_nodes[1].get_attribute("outputs:count")
        await og.Controller.evaluate()
        await og.Controller.evaluate()
        await og.Controller.evaluate()
        self.assertEqual(push_out_cnt_attr.get(), 3)
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(push_out_cnt_attr.get(), 3)
        for pipeline_stage in pipeline_stages:
            orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(pipeline_stage)
            self.assertEqual(len(orchestration_graphs), 1)
            orchestration_graph = orchestration_graphs[0]
            if pipeline_stage == og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND:
                self.assertEqual(len(orchestration_graph.get_nodes()), 1)
            else:
                self.assertEqual(len(orchestration_graph.get_nodes()), 0)

        # Confirm that newly added OnDemand graphs (after the initial population)
        # will properly trigger a repopulation of the OnDemand stage to include
        # those new graphs, thus allowing them to be executed. Also check that
        # the OnDemand orchestration graph now has two graph nodes, while the others
        # still have zero.
        exec_graph, exec_nodes = create_graph(
            "/ExecutionGraph", "execution", og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND
        )
        exec_out_cnt_attr = exec_nodes[1].get_attribute("outputs:count")
        await og.Controller.evaluate()
        await og.Controller.evaluate()
        await og.Controller.evaluate()
        self.assertEqual(push_out_cnt_attr.get(), 6)
        self.assertEqual(exec_out_cnt_attr.get(), 3)
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(push_out_cnt_attr.get(), 6)
        self.assertEqual(exec_out_cnt_attr.get(), 3)
        for pipeline_stage in pipeline_stages:
            orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(pipeline_stage)
            self.assertEqual(len(orchestration_graphs), 1)
            orchestration_graph = orchestration_graphs[0]
            if pipeline_stage == og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND:
                self.assertEqual(len(orchestration_graph.get_nodes()), 2)
            else:
                self.assertEqual(len(orchestration_graph.get_nodes()), 0)

        # Confirm that removing OnDemand graphs triggers a repopulation of the
        # OnDemand stage that won't include those graphs, resulting in them
        # no longer executing on an "on-demand" basis. Here we simply switch
        # the pipeline stage of the exec_graph to be something other than
        # OnDemand (Simulation). Check that both the Simulation and OnDemand
        # orchestration graphs now have one graph node apiece.
        og.cmds.ChangePipelineStage(
            graph=exec_graph, new_pipeline_stage=og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION
        )
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(push_out_cnt_attr.get(), 6)
        self.assertEqual(exec_out_cnt_attr.get(), 6)
        for pipeline_stage in pipeline_stages:
            orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(pipeline_stage)
            self.assertEqual(len(orchestration_graphs), 1)
            orchestration_graph = orchestration_graphs[0]
            if pipeline_stage in (
                og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND,
                og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION,
            ):
                self.assertEqual(len(orchestration_graph.get_nodes()), 1)
            else:
                self.assertEqual(len(orchestration_graph.get_nodes()), 0)

        # Confirm that editing OnDemand graphs triggers a repopulation of
        # the OnDemand stage, resulting in subsequent evaluations picking up
        # on the new changes. Here we check for this by creating an OnDemand
        # dirty_push graph and altering one of the underlying node connections
        # halfway through the evaluation checks; this change should trigger
        # a compute on the graph + should be picked up in the OnDemand pipeline.
        # Also check that the OnDemand orchestration graph now has two graph nodes,
        # while the Simulation orchestration graph still only has one.
        _, dirty_push_nodes = create_graph(
            "/DirtyPushGraph", "dirty_push", og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND
        )
        dirty_push_ontick_node = dirty_push_nodes[0]
        dirty_push_counter_node = dirty_push_nodes[1]
        dirty_push_out_cnt_attr = dirty_push_counter_node.get_attribute("outputs:count")
        await og.Controller.evaluate()
        await og.Controller.evaluate()
        await og.Controller.evaluate()
        self.assertEqual(push_out_cnt_attr.get(), 9)
        self.assertEqual(exec_out_cnt_attr.get(), 9)
        self.assertEqual(dirty_push_out_cnt_attr.get(), 1)
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(push_out_cnt_attr.get(), 9)
        self.assertEqual(exec_out_cnt_attr.get(), 12)
        self.assertEqual(dirty_push_out_cnt_attr.get(), 1)
        og.Controller.disconnect(
            dirty_push_ontick_node.get_attribute("outputs:tick"), dirty_push_counter_node.get_attribute("inputs:execIn")
        )
        await og.Controller.evaluate()
        await og.Controller.evaluate()
        await og.Controller.evaluate()
        self.assertEqual(push_out_cnt_attr.get(), 12)
        self.assertEqual(exec_out_cnt_attr.get(), 15)
        self.assertEqual(dirty_push_out_cnt_attr.get(), 2)
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(push_out_cnt_attr.get(), 12)
        self.assertEqual(exec_out_cnt_attr.get(), 18)
        self.assertEqual(dirty_push_out_cnt_attr.get(), 2)
        for pipeline_stage in pipeline_stages:
            orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(pipeline_stage)
            self.assertEqual(len(orchestration_graphs), 1)
            orchestration_graph = orchestration_graphs[0]
            if pipeline_stage == og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND:
                self.assertEqual(len(orchestration_graph.get_nodes()), 2)
            elif pipeline_stage == og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION:
                self.assertEqual(len(orchestration_graph.get_nodes()), 1)
            else:
                self.assertEqual(len(orchestration_graph.get_nodes()), 0)

    # ----------------------------------------------------------------------
    async def test_commands_with_group_undo_redo(self):
        """
        Tests direct and indirect commands work as expected when grouped under a single undo,
        including when nodes are created and destroyed
        """
        controller = og.Controller(update_usd=True)
        keys = controller.Keys
        (graph, (_, _, add3, add4), _, _) = controller.edit(
            "/World/Graph",
            {
                keys.CREATE_NODES: [
                    ("Add1", "omni.graph.nodes.Add"),
                    ("Add2", "omni.graph.nodes.Add"),
                    ("Add3", "omni.graph.nodes.Add"),
                    ("Add4", "omni.graph.nodes.Add"),
                ],
                keys.CONNECT: [
                    ("Add1.outputs:sum", "Add2.inputs:a"),
                ],
            },
        )

        og.cmds.DisableNode(node=add3)

        with omni.kit.undo.group():
            # add a new node
            controller.create_node("/World/Graph/Constant", "omni.graph.nodes.ConstantInt", undoable=True)
            node = graph.get_node("/World/Graph/Constant")
            controller.set(node.get_attribute("inputs:value"), 2, undoable=True)
            controller.connect("/World/Graph/Add1.outputs:sum", "/World/Graph/Add2.inputs:b", undoable=True)
            controller.disconnect("/World/Graph/Add1.outputs:sum", "/World/Graph/Add2.inputs:a", undoable=True)
            self.assertFalse(
                controller.attribute("/World/Graph/Add1.outputs:sum").is_connected(
                    controller.attribute("/World/Graph/Add2.inputs:a")
                )
            )
            og.cmds.EnableNode(node=add3)
            og.cmds.DisableNode(node=add4)
            controller.create_node("/World/Graph/Add5", "omni.graph.nodes.Add")
            og.cmds.DisableNode(node=graph.get_node("/World/Graph/Add5"))
            controller.delete_node("/World/Graph/Add2", update_usd=True, undoable=True)

        for _ in range(0, 3):
            graph = og.get_graph_by_path("/World/Graph")
            self.assertTrue(graph.is_valid())
            self.assertTrue(graph.get_node("/World/Graph/Constant").is_valid())
            self.assertTrue(graph.get_node("/World/Graph/Add1").is_valid())
            self.assertFalse(graph.get_node("/World/Graph/Add2").is_valid())
            self.assertEqual(controller.get(controller.attribute("/World/Graph/Constant.inputs:value")), 2)
            self.assertFalse(graph.get_node("/World/Graph/Add3").is_disabled())
            self.assertTrue(graph.get_node("/World/Graph/Add4").is_disabled())
            self.assertTrue(graph.get_node("/World/Graph/Add5").is_disabled())

            omni.kit.undo.undo()

            graph = og.get_graph_by_path("/World/Graph")
            self.assertTrue(graph.is_valid())
            self.assertFalse(graph.get_node("/World/Graph/Constant").is_valid())
            self.assertTrue(graph.get_node("/World/Graph/Add1").is_valid())
            self.assertTrue(graph.get_node("/World/Graph/Add2").is_valid())
            self.assertTrue(graph.get_node("/World/Graph/Add3").is_disabled())
            self.assertFalse(graph.get_node("/World/Graph/Add4").is_disabled())
            self.assertFalse(graph.get_node("/World/Graph/Add5").is_valid())

            omni.kit.undo.redo()

    # ----------------------------------------------------------------------
    async def test_create_graph_with_group_undo_redo(self):
        """
        Tests that commands perform as expected when the graph creation is part of the
        same undo group
        """
        orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(
            og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION
        )
        self.assertEqual(len(orchestration_graphs), 1)
        orchestration_graph = orchestration_graphs[0]

        with omni.kit.undo.group():
            (result, _) = og.cmds.CreateGraphAsNode(
                graph=orchestration_graph,
                node_name="my_push_graph",
                graph_path="/World/Graph",
                evaluator_name="push",
                is_global_graph=True,
                backed_by_usd=True,
                fc_backing_type=og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED,
                pipeline_stage=og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION,
                evaluation_mode=og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_STANDALONE,
            )
            self.assertTrue(result)
            graph = og.Controller.graph("/World/Graph")
            og.cmds.CreateNode(
                graph=graph,
                node_path="/World/Graph/ConstantInt",
                node_type="omni.graph.nodes.ConstantInt",
                create_usd=True,
            )
            og.cmds.SetAttr(
                attr=og.Controller.attribute("/World/Graph/ConstantInt.inputs:value"), value=5, update_usd=True
            )
            og.cmds.CreateNode(
                graph=graph, node_path="/World/Graph/Add", node_type="omni.graph.nodes.Add", create_usd=True
            )
            og.cmds.ConnectAttrs(
                src_attr="/World/Graph/ConstantInt.inputs:value", dest_attr="/World/Graph/Add.inputs:a", modify_usd=True
            )
            og.cmds.CreateVariable(
                graph=graph,
                variable_name="int_var",
                variable_type=og.Type(og.BaseDataType.INT),
                variable_value=5,
            )

        for _ in range(0, 3):
            self.assertTrue(og.Controller.graph("/World/Graph").is_valid())
            self.assertTrue(og.Controller.node("/World/Graph/ConstantInt").is_valid())
            self.assertTrue(og.Controller.node("/World/Graph/Add").is_valid())
            self.assertTrue(og.Controller.attribute("/World/Graph/ConstantInt.inputs:value").is_valid())
            self.assertTrue(og.Controller.attribute("/World/Graph/Add.inputs:a").is_valid())
            self.assertEquals(og.Controller.get(og.Controller.attribute("/World/Graph/ConstantInt.inputs:value")), 5)
            self.assertTrue(
                og.Controller.attribute("/World/Graph/Add.inputs:a").is_connected(
                    og.Controller.attribute("/World/Graph/ConstantInt.inputs:value")
                )
            )
            context = og.Controller.graph("/World/Graph").get_default_graph_context()
            self.assertEquals(og.Controller.variable("/World/Graph.graph:variable:int_var").get(context), 5)

            omni.kit.undo.undo()

            self.assertIsNone(og.Controller.graph("/World/Graph"))
            with self.assertRaises(og.OmniGraphError):
                _ = og.Controller.node("/World/Graph/ConstantInt")
            with self.assertRaises(og.OmniGraphError):
                _ = og.Controller.node("/World/Graph/Add")
            with self.assertRaises(og.OmniGraphError):
                _ = og.Controller.attribute("/World/Graph/ConstantInt.inputs:value")
            with self.assertRaises(og.OmniGraphError):
                _ = og.Controller.attribute("/World/Graph/Add.inputs:a")

            omni.kit.undo.redo()

    # ----------------------------------------------------------------------

    async def test_resolve_attr_type_command(self):
        """Tests ResolveAttrTypeCommand"""
        controller = og.Controller()
        keys = controller.Keys
        (_, (add1,), _, _) = controller.edit(
            "/World/Graph",
            {
                keys.CREATE_NODES: [
                    ("Add1", "omni.graph.nodes.Add"),
                ],
            },
        )
        # Test undo-redo for simple type
        unknown_t = controller.attribute_type("unknown")
        a = controller.attribute("inputs:a", add1)
        self.assertEqual(a.get_resolved_type(), unknown_t)
        og.cmds.ResolveAttrType(attr=a, type_id="int")
        self.assertEqual(a.get_resolved_type(), controller.attribute_type("int"))
        omni.kit.undo.undo()
        self.assertEqual(a.get_resolved_type(), unknown_t)
        omni.kit.undo.redo()
        self.assertEqual(a.get_resolved_type(), controller.attribute_type("int"))
