"""Action Graph Node Tests, Part 1"""

import time

import carb
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.app
import omni.kit.test
import omni.usd
from omni.graph.core import ThreadsafetyTestUtils
from pxr import Gf, OmniGraphSchemaTools, Sdf, Vt

# ======================================================================


class TestActionGraphNodes(ogts.OmniGraphTestCase):
    """Tests action graph node functionality"""

    TEST_GRAPH_PATH = "/World/TestGraph"
    keys = og.Controller.Keys
    E = og.ExecutionAttributeState.ENABLED
    D = og.ExecutionAttributeState.DISABLED

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

    # -----------------------------------------------------------------------
    async def test_addprimrelationship_node(self):
        """Test AddPrimRelationship node"""
        # Check that we can add a relationship to a prim and get that relationship.
        #
        # +---------+   +---------------------+
        # | ONTICK  +-->| AddPrimRelationship +
        # +---------+   +---------------------+
        #
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        og.Controller.create_graph({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})
        og.Controller.edit(
            self.TEST_GRAPH_PATH,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("AddRel", "omni.graph.action.AddPrimRelationship"),
                ],
                self.keys.CREATE_PRIMS: [
                    ("/Test", {}),
                    ("/Target", {}),
                ],
                self.keys.CONNECT: [("OnTick.outputs:tick", "AddRel.inputs:execIn")],
                self.keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("AddRel.inputs:path", "/Test"),
                    ("AddRel.inputs:name", "rel"),
                    ("AddRel.inputs:target", "/Target"),
                ],
            },
        )

        prim = stage.GetPrimAtPath("/Test")

        await og.Controller.evaluate()

        rel = prim.GetRelationship("rel")
        targets = rel.GetTargets()
        self.assertEqual(len(targets), 1)
        self.assertTrue(str(targets[0]) == "/Target")

    # ----------------------------------------------------------------------
    # The Branch node has a built-in test construct in its .ogn file located at ../../nodes/OgnBranch.ogn
    # (relative to the source location of the currently-opened testing script) AND is used in other testing
    # methods, so we skip adding extra node-specific tests for it here.

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_counter_node(self, test_instance_id: int = 0):
        """Test Counter node"""

        # Instance a test graph setup. Note that we append the graph path with the test_instance_id
        # so that the graph can be uniquely identified in the thread-safety test!
        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)
        og.Controller.create_graph({"graph_path": graph_path, "evaluator_name": "execution"})
        (_, (ontick_node, counter_node), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                self.keys.SET_VALUES: [("OnTick.inputs:onlyPlayback", False)],
                self.keys.CONNECT: ("OnTick.outputs:tick", "Counter.inputs:execIn"),
            },
        )

        # Obtain necessary attributes.
        in_exec_attr = counter_node.get_attribute("inputs:execIn")
        in_reset_attr = counter_node.get_attribute("inputs:reset")
        state_cnt_attr = counter_node.get_attribute("state:count")
        out_exec_attr = counter_node.get_attribute("outputs:execOut")
        out_cnt_attr = counter_node.get_attribute("outputs:count")
        out_tick_attr = ontick_node.get_attribute("outputs:tick")

        # Check that the counter node gets correctly incremented when executing.
        self.assertEqual(state_cnt_attr.get(), 0)
        self.assertEqual(out_cnt_attr.get(), 0)
        self.assertEqual(out_exec_attr.get(), self.D)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(state_cnt_attr.get(), 1)
        self.assertEqual(out_cnt_attr.get(), 1)
        self.assertEqual(out_exec_attr.get(), self.E)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(state_cnt_attr.get(), 2)
        self.assertEqual(out_cnt_attr.get(), 2)
        self.assertEqual(out_exec_attr.get(), self.E)

        # Check that the counter node doesn't increment when not executing.
        og.Controller.disconnect(
            out_tick_attr,
            in_exec_attr,
        )
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(state_cnt_attr.get(), 2)
        self.assertEqual(out_cnt_attr.get(), 2)
        self.assertEqual(out_exec_attr.get(), self.E)

        # Check that the reset flag for the Counter node instance works correctly when
        # inputs:execIn is set to 0 (i.e. when the Counter node is NOT supposed to be
        # executing).
        og.Controller.connect(out_tick_attr, in_reset_attr)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(state_cnt_attr.get(), 0)
        self.assertEqual(out_cnt_attr.get(), 0)
        self.assertEqual(out_exec_attr.get(), self.E)

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_delay_node(self, test_instance_id: int = 0):
        """Test Delay node"""

        # Instance a test graph setup. Note that we append the graph path with the test_instance_id
        # so that the graph can be uniquely identified in the thread-safety test!
        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)
        og.Controller.create_graph({"graph_path": graph_path, "evaluator_name": "execution"})
        (_, (on_impulse_node, _, counter_node), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("Delay", "omni.graph.action.Delay"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                self.keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "Delay.inputs:execIn"),
                    ("Delay.outputs:finished", "Counter.inputs:execIn"),
                ],
                self.keys.SET_VALUES: [
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("Delay.inputs:duration", 0.01),
                ],
            },
        )

        # Obtain necessary attributes.
        out_cnt_attr = counter_node.get_attribute("outputs:count")
        state_enable_impulse_attr = on_impulse_node.get_attribute("state:enableImpulse")

        # Trigger the graph(s) once.
        state_enable_impulse_attr.set(True)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.

        # Downstream execution is delayed, so the counter node won't be incremented.
        self.assertEqual(out_cnt_attr.get(), 0)

        # Wait to ensure that the delay finishes before checking that the counter node
        # has indeed been incremented.
        time.sleep(0.02)
        yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
        self.assertEqual(out_cnt_attr.get(), 1)

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_flipflop_node(self, test_instance_id: int = 0):
        """Test FlipFlop node"""

        # Instance a test graph setup. Note that we append the graph path with the test_instance_id
        # so that the graph can be uniquely identified in the thread-safety test!
        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)
        og.Controller.create_graph({"graph_path": graph_path, "evaluator_name": "execution"})
        (graph, (_, flip_flop_node), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("FlipFlop", "omni.graph.action.FlipFlop"),
                ],
                self.keys.SET_VALUES: [("OnTick.inputs:onlyPlayback", False)],
                self.keys.CONNECT: (
                    "OnTick.outputs:tick",
                    "FlipFlop.inputs:execIn",
                ),
            },
        )

        # Obtain necessary attributes.
        out_a_attr = flip_flop_node.get_attribute("outputs:a")
        out_b_attr = flip_flop_node.get_attribute("outputs:b")
        out_isa_attr = flip_flop_node.get_attribute("outputs:isA")

        # First eval, 'a'.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_a_attr.get(), self.E)
        self.assertEqual(out_b_attr.get(), self.D)
        self.assertTrue(out_isa_attr.get())

        # Second eval 'b'.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_a_attr.get(), self.D)
        self.assertEqual(out_b_attr.get(), self.E)
        self.assertFalse(out_isa_attr.get())

        # Third eval 'a'.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_a_attr.get(), self.E)
        self.assertEqual(out_b_attr.get(), self.D)
        self.assertTrue(out_isa_attr.get())

        # Test that non-usd-backed FlipFlop nodes correctly
        # set their execution-type attributes.

        # Make sure that the node paths are prefaced with the path to the
        # graph that they should reside in (so that each node creation command
        # can be processed to produce unique nodes for each test instance)!
        _, on_tick_no_usd = og.cmds.CreateNode(
            graph=graph,
            node_path=f"{graph_path}/OnTickNoUSD",
            node_type="omni.graph.action.OnTick",
            create_usd=False,
        )
        _, flip_flop_no_usd = og.cmds.CreateNode(
            graph=graph,
            node_path=f"{graph_path}/FlipFlopNoUSD",
            node_type="omni.graph.action.FlipFlop",
            create_usd=False,
        )

        # Obtain necessary attributes.
        out_a_attr = flip_flop_no_usd.get_attribute("outputs:a")
        out_b_attr = flip_flop_no_usd.get_attribute("outputs:b")
        out_isa_attr = flip_flop_no_usd.get_attribute("outputs:isA")

        on_tick_no_usd.get_attribute("inputs:onlyPlayback").set(False)

        # Make sure that the node attribute paths are prefaced with the graph
        # path that they reside in (so that each instanced node attribute can
        # be uniquely processed)!
        og.cmds.ConnectAttrs(
            src_attr=f"{graph_path}/OnTickNoUSD.outputs:tick",
            dest_attr=f"{graph_path}/FlipFlopNoUSD.inputs:execIn",
            modify_usd=False,
        )

        # First eval, 'a'.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_a_attr.get(), self.E)
        self.assertEqual(out_b_attr.get(), self.D)
        self.assertTrue(out_isa_attr.get())

        # Second eval 'b'.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_a_attr.get(), self.D)
        self.assertEqual(out_b_attr.get(), self.E)
        self.assertFalse(out_isa_attr.get())

        # Third eval 'a'.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_a_attr.get(), self.E)
        self.assertEqual(out_b_attr.get(), self.D)
        self.assertTrue(out_isa_attr.get())

    # ----------------------------------------------------------------------
    # The ForEach node has a built-in test construct in its .ogn file located at ../../nodes/OgnForEach.ogn
    # (relative to the source location of the currently-opened testing script), so we skip adding extra
    # node-specific tests for it here.

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

        # Since we want to use the same prim across all graph instances in the
        # thread-safety test, we add it to the threading cache like so:
        prim = ThreadsafetyTestUtils.add_to_threading_cache(test_instance_id, stage.DefinePrim("/World/TestPrim"))
        ThreadsafetyTestUtils.add_to_threading_cache(
            test_instance_id,
            prim.CreateAttribute("val1", Sdf.ValueTypeNames.Int2, False).Set(Gf.Vec2i(1, 1)),
        )

        # Instance a test graph setup. Note that we append the graph path with the test_instance_id
        # so that the graph can be uniquely identified in the thread-safety test!
        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)
        og.Controller.create_graph({"graph_path": graph_path, "evaluator_name": "execution"})
        (_, (_, _, _, _, _, _, write_node, finish_counter), _, _) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Const", "omni.graph.nodes.ConstantInt2"),
                    ("StopNum", "omni.graph.nodes.ConstantInt"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("Branch", "omni.graph.action.Branch"),
                    ("For", "omni.graph.action.ForLoop"),
                    ("Write1", "omni.graph.nodes.WritePrimAttribute"),
                    ("FinishCounter", "omni.graph.action.Counter"),
                ],
                self.keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("Const.inputs:value", [1, 2]),
                    ("StopNum.inputs:value", 3),
                    ("Write1.inputs:name", "val1"),
                    ("Write1.inputs:primPath", "/World/TestPrim"),
                    ("Write1.inputs:usePath", True),
                    ("Branch.inputs:condition", True),
                ],
                self.keys.CONNECT: [
                    ("OnTick.outputs:tick", "For.inputs:execIn"),
                    ("StopNum.inputs:value", "For.inputs:stop"),
                    ("For.outputs:loopBody", "Branch.inputs:execIn"),
                    ("For.outputs:finished", "FinishCounter.inputs:execIn"),
                    ("Branch.outputs:execTrue", "Write1.inputs:execIn"),
                    ("For.outputs:value", "Add.inputs:a"),
                    ("Const.inputs:value", "Add.inputs:b"),
                    ("Add.outputs:sum", "Write1.inputs:value"),
                ],
            },
        )

        # Evaluate the graph(s).
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertListEqual([3, 4], list(stage.GetAttributeAtPath("/World/TestPrim.val1").Get()))
        self.assertEqual(3, write_node.get_compute_count())
        self.assertEqual(1, finish_counter.get_compute_count())

    # This tests make sure that a for loop writing arrays to diferent prims
    # work as expected (OM-84129)
    async def test_foreach_node_write_multiple_prim(self, test_instance_id: int = 0):
        """Test the foreach node writing arrays to output prims"""
        og.Controller.create_graph({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})
        (_, _, (prim1, prim2), _) = og.Controller.edit(
            self.TEST_GRAPH_PATH,
            {
                self.keys.CREATE_PRIMS: [
                    ("/World/Prim1", {"graph_output": ("Int[]", [])}),
                    ("/World/Prim2", {"graph_output": ("Int[]", [])}),
                ],
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("For", "omni.graph.action.ForEach"),
                    ("Write", "omni.graph.nodes.WritePrimAttribute"),
                    ("MakeArray", "omni.graph.nodes.ConstructArray"),
                ],
                self.keys.SET_VALUES: [
                    ("For.inputs:arrayIn", {"type": "token[]", "value": ["/World/Prim1", "/World/Prim2"]}),
                    ("OnTick.inputs:onlyPlayback", False),
                    ("Write.inputs:name", "graph_output"),
                    ("Write.inputs:usePath", True),
                    ("Write.inputs:usdWriteBack", True),
                    ("MakeArray.inputs:arraySize", 1),
                ],
                self.keys.CONNECT: [
                    ("OnTick.outputs:tick", "For.inputs:execIn"),
                    ("For.outputs:loopBody", "Write.inputs:execIn"),
                    ("For.outputs:element", "Write.inputs:primPath"),
                    ("For.outputs:arrayIndex", "MakeArray.inputs:input0"),
                    ("MakeArray.outputs:array", "Write.inputs:value"),
                ],
            },
        )

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

        prim1_out = prim1.GetAttribute("graph_output")
        prim2_out = prim2.GetAttribute("graph_output")

        self.assertEqual(prim1_out.Get(), Vt.IntArray([0]))
        self.assertEqual(prim2_out.Get(), Vt.IntArray([1]))

    # ----------------------------------------------------------------------
    # The Gate node has a built-in test construct in its .ogn file located at ../../nodes/OgnGate.ogn
    # (relative to the source location of the currently-opened testing script) AND is used in other
    # testing methods, so we skip adding extra node-specific tests for it here.

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_multigate_node(self, test_instance_id: int = 0):
        """Test Multigate node"""

        # Instance a test graph setup. Note that we append the graph path with the test_instance_id
        # so that the graph can be uniquely identified in the thread-safety test!
        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)
        (_, (_, multigate_node), _, _) = og.Controller.edit(
            {"graph_path": graph_path, "evaluator_name": "execution"},
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Multigate", "omni.graph.action.Multigate"),
                ],
                self.keys.SET_VALUES: [("OnTick.inputs:onlyPlayback", False)],
                self.keys.CONNECT: [
                    ("OnTick.outputs:tick", "Multigate.inputs:execIn"),
                ],
            },
        )

        # Add 5 extra outputs to the Multigate node.
        for i in range(1, 6):
            og.Controller.create_attribute(
                multigate_node,
                f"outputs:output{i}",
                og.Type(og.BaseDataType.UINT, 1, 0, og.AttributeRole.EXECUTION),
                og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT,
            )

        # Obtain necessary attributes.
        out_attr_0 = multigate_node.get_attribute("outputs:output0")
        out_attr_1 = multigate_node.get_attribute("outputs:output1")
        out_attr_2 = multigate_node.get_attribute("outputs:output2")
        out_attr_3 = multigate_node.get_attribute("outputs:output3")
        out_attr_4 = multigate_node.get_attribute("outputs:output4")
        out_attr_5 = multigate_node.get_attribute("outputs:output5")

        # Check that the Multigate node correctly cycles through each of its outputs.
        # Note that we trigger an execution through the Multigate node via the OnTick node,
        # whose onlyPlayback input we've set to False in order to trigger an execution each
        # time we evaluate the graph(s).
        for i in range(0, 6):
            yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
            self.assertEqual(multigate_node.get_attribute(f"outputs:output{(i + 0) % 6}").get(), self.E)
            self.assertEqual(multigate_node.get_attribute(f"outputs:output{(i + 1) % 6}").get(), self.D)
            self.assertEqual(multigate_node.get_attribute(f"outputs:output{(i + 2) % 6}").get(), self.D)
            self.assertEqual(multigate_node.get_attribute(f"outputs:output{(i + 3) % 6}").get(), self.D)
            self.assertEqual(multigate_node.get_attribute(f"outputs:output{(i + 4) % 6}").get(), self.D)
            self.assertEqual(multigate_node.get_attribute(f"outputs:output{(i + 5) % 6}").get(), self.D)

        # Next try removing some output attributes during evaluation and test if the
        # Multigate node correctly cycles through.
        for _ in range(4):
            yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_attr_0.get(), self.D)
        self.assertEqual(out_attr_1.get(), self.D)
        self.assertEqual(out_attr_2.get(), self.D)
        self.assertEqual(out_attr_3.get(), self.E)
        self.assertEqual(out_attr_4.get(), self.D)
        self.assertEqual(out_attr_5.get(), self.D)
        multigate_node.remove_attribute("outputs:output4")

        # The Multigate node cycles back to 0 instead of going to 5 since it thinks that it's
        # reached the end of its outputs list (i.e. it expects to jump from pin 3 to 4, but b/c
        # there is no such pin it goes back to 0 rather than 5).
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_attr_0.get(), self.E)
        self.assertEqual(out_attr_1.get(), self.D)
        self.assertEqual(out_attr_2.get(), self.D)
        self.assertEqual(out_attr_3.get(), self.D)
        self.assertEqual(out_attr_5.get(), self.D)

        # Further showing that executing 4 times brings us back to pin 0 rather than pin 5.
        for _ in range(4):
            yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_attr_0.get(), self.E)
        self.assertEqual(out_attr_1.get(), self.D)
        self.assertEqual(out_attr_2.get(), self.D)
        self.assertEqual(out_attr_3.get(), self.D)
        self.assertEqual(out_attr_5.get(), self.D)

        # Execute the graph(s) once, then remove the currently-enabled output pin. The Multigate node will think
        # that it's reached the end of the outputs list, and cycle back. Because we removed pin 1, it'll go
        # back to pin 0 and never cycle through the other outputs since it cannot make the jump from
        # pin 0 to 2 (as mentioned previously).
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        multigate_node.remove_attribute("outputs:output1")
        for _ in range(3):
            yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_attr_0.get(), self.E)
        self.assertEqual(out_attr_2.get(), self.D)
        self.assertEqual(out_attr_3.get(), self.D)
        self.assertEqual(out_attr_5.get(), self.D)

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_multisequence_node(self, test_instance_id: int = 0):
        """Test Multisequence node"""

        # Instance a test graph setup. Note that we append the graph path with the test_instance_id
        # so that the graph can be uniquely identified in the thread-safety test!
        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)
        og.Controller.create_graph({"graph_path": graph_path, "evaluator_name": "execution"})
        (
            _,
            (on_tick_node, multisequence_node, counter0_node, counter1_node, counter2_node),
            _,
            _,
        ) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Multisequence", "omni.graph.action.Multisequence"),
                    ("Counter0", "omni.graph.action.Counter"),
                    ("Counter1", "omni.graph.action.Counter"),
                    ("Counter2", "omni.graph.action.Counter"),
                ],
                self.keys.SET_VALUES: [("OnTick.inputs:onlyPlayback", False)],
                self.keys.CONNECT: [
                    ("OnTick.outputs:tick", "Multisequence.inputs:execIn"),
                ],
            },
        )

        # Add 3 extra outputs to the Multisequence node.
        for j in range(1, 4):
            og.Controller.create_attribute(
                multisequence_node,
                f"outputs:output{j}",
                og.Type(og.BaseDataType.UINT, 1, 0, og.AttributeRole.EXECUTION),
                og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT,
            )

        # Obtain necessary attributes.
        out_attr_0 = multisequence_node.get_attribute("outputs:output0")
        out_attr_1 = multisequence_node.get_attribute("outputs:output1")
        out_attr_2 = multisequence_node.get_attribute("outputs:output2")
        out_attr_3 = multisequence_node.get_attribute("outputs:output3")
        in_exec_attr_0 = counter0_node.get_attribute("inputs:execIn")
        in_exec_attr_1 = counter1_node.get_attribute("inputs:execIn")
        in_exec_attr_2 = counter2_node.get_attribute("inputs:execIn")
        out_cnt_attr_0 = counter0_node.get_attribute("outputs:count")
        out_cnt_attr_1 = counter1_node.get_attribute("outputs:count")
        out_cnt_attr_2 = counter2_node.get_attribute("outputs:count")
        in_onlyplayback_attr = on_tick_node.get_attribute("inputs:onlyPlayback")

        # Connect Multisequence node output attributes to the counter nodes.
        og.Controller.connect(out_attr_0, in_exec_attr_0)
        og.Controller.connect(out_attr_2, in_exec_attr_1)
        og.Controller.connect(out_attr_1, in_exec_attr_2)

        # Check that the Multisequence node correctly executes through its outputs when input
        # execution is enabled via the OnTick node. This is done by checking whether
        # each counter has been incremented by 1, and if the last output pin on the
        # Multisequence node remains enabled (regardless of the fact that it's not connected
        # downstream).
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_cnt_attr_0.get(), 1)
        self.assertEqual(out_cnt_attr_1.get(), 1)
        self.assertEqual(out_cnt_attr_2.get(), 1)
        self.assertEqual(out_attr_0.get(), self.D)
        self.assertEqual(out_attr_1.get(), self.D)
        self.assertEqual(out_attr_2.get(), self.D)
        self.assertEqual(out_attr_3.get(), self.E)

        # Connect the Counter2 node to another Multisequence output pin.
        og.Controller.connect(out_attr_3, in_exec_attr_2)

        # Once again evaluate the graph(s). In this situation the Counter2 node should be incremented twice
        # (since it's connected to 2 separate Multisequence output pins). Also Multisequence output pins
        # 1 AND 3 should both be enabled by the end of the execution; this is because pin 3 would
        # typically be the last output that gets enabled, but because pin 3 shares a downstream
        # node with pin 1 (that node being Counter2), both outputs need to be enabled by the end.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_cnt_attr_0.get(), 2)
        self.assertEqual(out_cnt_attr_1.get(), 2)
        self.assertEqual(out_cnt_attr_2.get(), 3)
        self.assertEqual(out_attr_0.get(), self.D)
        self.assertEqual(out_attr_1.get(), self.E)
        self.assertEqual(out_attr_2.get(), self.D)
        self.assertEqual(out_attr_3.get(), self.E)

        # Set the OnTick node to only trigger downstream execution when playback is enabled; check
        # that in this situation the Multisequence node correctly skips executing through its outputs
        # (i.e. that the Counter nodes don't get incremented). The state of the Multisequence's output
        # pins should not have changed since the last graph evaluation.
        in_onlyplayback_attr.set(True)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_cnt_attr_0.get(), 2)
        self.assertEqual(out_cnt_attr_1.get(), 2)
        self.assertEqual(out_cnt_attr_2.get(), 3)
        self.assertEqual(out_attr_0.get(), self.D)
        self.assertEqual(out_attr_1.get(), self.E)
        self.assertEqual(out_attr_2.get(), self.D)
        self.assertEqual(out_attr_3.get(), self.E)

        # Try removing an output attribute from the Multisequence node, check that all other
        # outputs that come before in the list get triggered. In this example we remove outputs2,
        # which means that outputs3 won't get triggered at all (as evidenced by the fact that the
        # Counter2 node only gets incremented once by output1).
        og.Controller.disconnect(out_attr_2, in_exec_attr_1)
        multisequence_node.remove_attribute("outputs:output2")
        in_onlyplayback_attr.set(False)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_cnt_attr_0.get(), 3)
        self.assertEqual(out_cnt_attr_1.get(), 2)
        self.assertEqual(out_cnt_attr_2.get(), 4)
        self.assertEqual(out_attr_0.get(), self.D)
        self.assertEqual(out_attr_1.get(), self.E)
        self.assertEqual(out_attr_3.get(), self.D)

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_once_node(self, test_instance_id: int = 0):
        """Test Once node"""

        # Instance a test graph setup. Note that we append the graph path with the test_instance_id
        # so that the graph can be uniquely identified in the thread-safety test!
        graph_path = self.TEST_GRAPH_PATH + str(test_instance_id)
        og.Controller.create_graph({"graph_path": graph_path, "evaluator_name": "execution"})
        (_, (ontick_node, once_node), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Once", "omni.graph.action.Once"),
                ],
                self.keys.SET_VALUES: [("OnTick.inputs:onlyPlayback", False)],
                self.keys.CONNECT: ("OnTick.outputs:tick", "Once.inputs:execIn"),
            },
        )

        # Obtain necessary attributes.
        in_exec_attr = once_node.get_attribute("inputs:execIn")
        in_reset_attr = once_node.get_attribute("inputs:reset")
        out_once_attr = once_node.get_attribute("outputs:once")
        out_after_attr = once_node.get_attribute("outputs:after")
        out_tick_attr = ontick_node.get_attribute("outputs:tick")

        # Check that the Once node controls flow of execution by passing flow
        # differently the first time it's executed compared to all subsequent
        # executions.
        self.assertEqual(out_once_attr.get(), self.D)
        self.assertEqual(out_after_attr.get(), self.D)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_once_attr.get(), self.E)
        self.assertEqual(out_after_attr.get(), self.D)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_once_attr.get(), self.D)
        self.assertEqual(out_after_attr.get(), self.E)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_once_attr.get(), self.D)
        self.assertEqual(out_after_attr.get(), self.E)

        # Check that the reset flag works correctly when inputs:execIn is set to 0 (i.e. when
        # the Once node is NOT supposed to be executing).
        og.Controller.disconnect(out_tick_attr, in_exec_attr)
        og.Controller.connect(out_tick_attr, in_reset_attr)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_once_attr.get(), self.D)
        self.assertEqual(out_after_attr.get(), self.D)
        og.Controller.disconnect(out_tick_attr, in_reset_attr)
        og.Controller.connect(out_tick_attr, in_exec_attr)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_once_attr.get(), self.E)
        self.assertEqual(out_after_attr.get(), self.D)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_once_attr.get(), self.D)
        self.assertEqual(out_after_attr.get(), self.E)

        # Check that when both the execIn and reset input attributes get triggered, the latter
        # overrides the former and execution flow does not pass through outputs:after.
        # FIXME: Something about the 2nd connection here is messing up the data model such
        # that inputs:reset is being read as 0 inside the node
        og.Controller.connect(out_tick_attr, in_reset_attr)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        # self.assertEqual(out_once_attr.get(), self.D)
        # self.assertEqual(out_after_attr.get(), self.D)

    # ----------------------------------------------------------------------
    # NOTE: Even though the OnClosing node is threadsafe (its compute method is very simple),
    # we don't adapt the below test to check for thread-safety conditions because it relies
    # on other nodes (omni.graph.action.SendCustomEvent and omni.graph.nodes.GraphTarget)
    # which are NOT threadsafe.
    async def test_onclosing_node(self):
        """Test OnClosing node"""
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})

        # Testing OnClosing is tricky because OG is being destroyed when it happens -
        # so test by sending a custom event when the network is triggered
        # and then checking if we got that event.

        def registered_event_name(event_name):
            """Returns the internal name used for the given custom event name"""
            name = "omni.graph.action." + event_name
            return carb.events.type_from_string(name)

        got_event = [0]

        def on_event(_):
            got_event[0] = got_event[0] + 1

        reg_event_name = registered_event_name("foo")
        message_bus = omni.kit.app.get_app().get_message_bus_event_stream()
        sub = message_bus.create_subscription_to_push_by_type(reg_event_name, on_event)
        self.assertIsNotNone(sub)

        async def set_up_graph():
            controller = og.Controller()
            keys = og.Controller.Keys
            controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.CREATE_NODES: [
                        ("OnClosing", "omni.graph.action.OnClosing"),
                        ("Send", "omni.graph.action.SendCustomEvent"),
                        ("GraphTarget", "omni.graph.nodes.GraphTarget"),
                    ],
                    keys.CONNECT: [
                        ("OnClosing.outputs:execOut", "Send.inputs:execIn"),
                        ("GraphTarget.outputs:primPath", "Send.inputs:path"),
                    ],
                    keys.SET_VALUES: [("Send.inputs:eventName", "foo")],
                },
            )

        await set_up_graph()
        # Evaluate once so that the standalone graph is in steady state.
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(got_event[0], 0)

        # Close the stage.
        usd_context = omni.usd.get_context()
        (result, _) = await usd_context.close_stage_async()
        self.assertTrue(result)
        # Check our handler was called.
        self.assertEqual(got_event[0], 1)
        # Reset the counter.
        got_event[0] = 0

        # Now check that the same works with instanced graphs.
        await usd_context.new_stage_async()
        await set_up_graph()

        og.cmds.SetEvaluationMode(
            graph=og.get_graph_by_path(self.TEST_GRAPH_PATH),
            new_evaluation_mode=og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_INSTANCED,
        )

        stage = usd_context.get_stage()

        prims = [stage.DefinePrim(f"/prim_{i}") for i in range(0, 100)]
        for prim in prims:
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim.GetPath(), self.TEST_GRAPH_PATH)

        # Wait an update for the graphs to get set up.
        await omni.kit.app.get_app().next_update_async()

        # Close the stage.
        (result, _) = await usd_context.close_stage_async()
        self.assertTrue(result)
        # Check that our handler was called.
        self.assertEqual(got_event[0], len(prims))

    # ----------------------------------------------------------------------
    async def test_oncustomevent_and_sendcustomevent_nodes(self):
        """Test OnCustomEvent and SendCustomEvent nodes"""
        controller = og.Controller()

        controller.create_graph({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})
        (_, (_, _, event1_node, counter1_node, event2_node, counter2_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                self.keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("Send", "omni.graph.action.SendCustomEvent"),
                    ("OnCustomEvent1", "omni.graph.action.OnCustomEvent"),
                    ("Counter1", "omni.graph.action.Counter"),
                    ("OnCustomEvent2", "omni.graph.action.OnCustomEvent"),
                    ("Counter2", "omni.graph.action.Counter"),
                ],
                self.keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "Send.inputs:execIn"),
                    ("OnCustomEvent1.outputs:execOut", "Counter1.inputs:execIn"),
                    ("OnCustomEvent2.outputs:execOut", "Counter2.inputs:execIn"),
                ],
                self.keys.SET_VALUES: [
                    ("OnCustomEvent1.inputs:onlyPlayback", False),
                    ("OnCustomEvent2.inputs:onlyPlayback", False),
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("Send.inputs:eventName", "foo"),
                    ("Send.inputs:path", "Test Path"),
                    ("OnCustomEvent1.inputs:eventName", "foo"),
                    ("OnCustomEvent2.inputs:eventName", "foo"),
                ],
            },
        )

        counter1_controller = og.Controller(og.Controller.attribute("outputs:count", counter1_node))
        counter2_controller = og.Controller(og.Controller.attribute("outputs:count", counter2_node))
        event1_controller = og.Controller(og.Controller.attribute("outputs:path", event1_node))
        event2_controller = og.Controller(og.Controller.attribute("outputs:path", event2_node))

        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(counter1_controller.get(), 0)
        # Trigger graph once, this will queue up the event for the next evaluation.
        controller.edit(self.TEST_GRAPH_PATH, {self.keys.SET_VALUES: ("OnImpulse.state:enableImpulse", True)})
        # Note that if this is a push subscription, the receivers will run this frame instead of next.
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(counter1_controller.get(), 0)
        # This evaluation should trigger the receivers.
        await omni.kit.app.get_app().next_update_async()
        # Verify that events were received.
        self.assertEqual(counter1_controller.get(), 1)
        self.assertEqual(event1_controller.get(), "Test Path")
        self.assertEqual(counter2_controller.get(), 1)
        self.assertEqual(event2_controller.get(), "Test Path")
        # Verify the contents of the associated bundle.
        # FIXME: Authored bundle is always empty?
        # bundle_contents = og.BundleContents(graph.get_default_graph_context(), event1_node, "outputs:bundle", True)
        # self.assertEqual(1, bundle_contents.size)

        # Modify the event name one receiver and sender and ensure it still works.
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                self.keys.SET_VALUES: [("Send.inputs:eventName", "bar"), ("OnImpulse.state:enableImpulse", True)],
            },
        )
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        # We changed the sender event name, so counter should NOT have triggered again.
        self.assertEqual(counter1_controller.get(), 1)
        # Change the receiver name to match.
        controller.edit(self.TEST_GRAPH_PATH, {self.keys.SET_VALUES: ("OnCustomEvent1.inputs:eventName", "bar")})
        await omni.kit.app.get_app().next_update_async()
        # Trigger send again and verify we get it (1 frame lag for pop).
        controller.edit(self.TEST_GRAPH_PATH, {self.keys.SET_VALUES: ("OnImpulse.state:enableImpulse", True)})
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(counter1_controller.get(), 2)
