"""Action Graph Node Tests, Part 3"""

import random
from functools import partial

import carb.input
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 carb.input import MouseEventType
from omni.graph.core import ThreadsafetyTestUtils
from pxr import Gf, OmniGraphSchemaTools, Sdf


# ======================================================================
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()

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_onmouseinput_node(self, test_instance_id: int = 0):
        """Test OnMouseInput 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_mouse_input_node,), _, _) = og.Controller.edit(
            graph_path,
            {self.keys.CREATE_NODES: ("OnMouseInput", "omni.graph.action.OnMouseInput")},
        )

        # Obtain necessary attributes.
        in_onlyplayback_attr = on_mouse_input_node.get_attribute("inputs:onlyPlayback")
        in_mouse_element_attr = on_mouse_input_node.get_attribute("inputs:mouseElement")
        out_pressed_attr = on_mouse_input_node.get_attribute("outputs:pressed")
        out_released_attr = on_mouse_input_node.get_attribute("outputs:released")
        out_valuechanged_attr = on_mouse_input_node.get_attribute("outputs:valueChanged")
        out_ispressed_attr = on_mouse_input_node.get_attribute("outputs:isPressed")
        out_value_attr = on_mouse_input_node.get_attribute("outputs:value")

        # Obtain an interface to the mouse and carb input provider.
        mouse = ThreadsafetyTestUtils.add_to_threading_cache(
            test_instance_id, omni.appwindow.get_default_app_window().get_mouse()
        )
        input_provider = ThreadsafetyTestUtils.add_to_threading_cache(
            test_instance_id, carb.input.acquire_input_provider()
        )

        # Define a list of tokens representing the possible inputs to the OnMouseInput node's "mouseElement" attribute.
        mouse_tokens = ThreadsafetyTestUtils.add_to_threading_cache(
            test_instance_id,
            ["Left Button", "Middle Button", "Right Button", "Normalized Move", "Pixel Move", "Scroll"],
        )

        # Define a list of all possible mouse event types (for convenience).
        mouse_event_types = ThreadsafetyTestUtils.add_to_threading_cache(
            test_instance_id,
            [
                MouseEventType.LEFT_BUTTON_DOWN,
                MouseEventType.LEFT_BUTTON_UP,
                MouseEventType.MIDDLE_BUTTON_DOWN,
                MouseEventType.MIDDLE_BUTTON_UP,
                MouseEventType.RIGHT_BUTTON_DOWN,
                MouseEventType.RIGHT_BUTTON_UP,
                MouseEventType.MOVE,
                MouseEventType.SCROLL,
            ],
        )

        # Define some imaginary window dimensions for testing purposes.
        window_width = ThreadsafetyTestUtils.add_to_threading_cache(test_instance_id, 1440)
        window_height = ThreadsafetyTestUtils.add_to_threading_cache(test_instance_id, 720)

        # Wrap main test driver code in the following generator (to leverage the
        # ThreadsafetyTestUtils.make_threading_test decorator to its fullest). Note that
        # in addition to testing whether specific buttons activate their corresponding code
        # path, we also check that no other buttons can activate code paths they don't belong
        # to (e.g. that pressing the left mouse button doesn't get registered as a button
        # press when "inputs:mouseElement" is set to "Right Button").
        def _test_onmouseinput_node():
            # Loop through each possible "inputs.mouseElement" token, and set the input attribute
            #  on the OnMouseInput nodes in each graph.
            for token in mouse_tokens:
                in_mouse_element_attr.set(token)
                # Loop through each possible mouse event type.
                for event_type in mouse_event_types:
                    # Generate a new random test position for the mouse cursor. Note that we
                    # add it to the threading cache so that each test graph instance runs
                    # its tests/comparisons against the same randomly-generated position.
                    pos = ThreadsafetyTestUtils.add_to_threading_cache(
                        test_instance_id, (random.randint(0, window_width), random.randint(0, window_height))
                    )
                    # Trigger the current mouse event.
                    input_provider.buffer_mouse_event(
                        mouse, event_type, (pos[0] / window_width, pos[1] / window_height), 0, pos
                    )
                    yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
                    # Test left/middle/right mouse button pressed functionality. Only left/middle/right
                    # mouse button presses should activate this codepath.
                    if (
                        (event_type == MouseEventType.LEFT_BUTTON_DOWN and token == "Left Button")
                        or (event_type == MouseEventType.MIDDLE_BUTTON_DOWN and token == "Middle Button")
                        or (event_type == MouseEventType.RIGHT_BUTTON_DOWN and token == "Right Button")
                    ):
                        self.assertEqual(out_pressed_attr.get(), self.E)
                        self.assertEqual(out_released_attr.get(), self.D)
                        self.assertEqual(out_valuechanged_attr.get(), self.D)
                        self.assertTrue(out_ispressed_attr.get())
                        self.assertEqual(len(out_value_attr.get()), 2)
                        self.assertEqual(out_value_attr.get()[0], 0)
                        self.assertEqual(out_value_attr.get()[1], 0)
                    # Test left/middle/right mouse button released functionality. Only left/middle/right mouse
                    # button releases should activate this codepath.
                    elif (
                        (event_type == MouseEventType.LEFT_BUTTON_UP and token == "Left Button")
                        or (event_type == MouseEventType.MIDDLE_BUTTON_UP and token == "Middle Button")
                        or (event_type == MouseEventType.RIGHT_BUTTON_UP and token == "Right Button")
                    ):
                        self.assertEqual(out_pressed_attr.get(), self.D)
                        self.assertEqual(out_released_attr.get(), self.E)
                        self.assertEqual(out_valuechanged_attr.get(), self.D)
                        self.assertFalse(out_ispressed_attr.get())
                        self.assertEqual(len(out_value_attr.get()), 2)
                        self.assertEqual(out_value_attr.get()[0], 0)
                        self.assertEqual(out_value_attr.get()[1], 0)
                    # Test mouse movement functionality with the "Normalized Move" option enabled. Only mouse movement
                    # should activate this codepath.
                    elif event_type == MouseEventType.MOVE and token == "Normalized Move":
                        self.assertEqual(out_pressed_attr.get(), self.D)
                        self.assertEqual(out_released_attr.get(), self.D)
                        self.assertEqual(out_valuechanged_attr.get(), self.E)
                        self.assertFalse(out_ispressed_attr.get())
                        self.assertEqual(len(out_value_attr.get()), 2)
                        self.assertAlmostEqual(out_value_attr.get()[0], pos[0] / window_width)
                        self.assertAlmostEqual(out_value_attr.get()[1], pos[1] / window_height)
                    # Test mouse movement functionality with the "Pixel Move" option enabled. Only mouse movement
                    # should activate this codepath.
                    elif event_type == MouseEventType.MOVE and token == "Pixel Move":
                        self.assertEqual(out_pressed_attr.get(), self.D)
                        self.assertEqual(out_released_attr.get(), self.D)
                        self.assertEqual(out_valuechanged_attr.get(), self.E)
                        self.assertFalse(out_ispressed_attr.get())
                        self.assertEqual(len(out_value_attr.get()), 2)
                        self.assertEqual(out_value_attr.get()[0], pos[0])
                        self.assertEqual(out_value_attr.get()[1], pos[1])
                    # Test mouse scrolling functionality with the "Scroll" option enabled. Only mouse scrolling
                    # should activate this codepath.
                    elif event_type == MouseEventType.SCROLL and token == "Scroll":
                        self.assertEqual(out_pressed_attr.get(), self.D)
                        self.assertEqual(out_released_attr.get(), self.D)
                        self.assertEqual(out_valuechanged_attr.get(), self.E)
                        self.assertFalse(out_ispressed_attr.get())
                        self.assertEqual(len(out_value_attr.get()), 2)
                        self.assertAlmostEqual(out_value_attr.get()[0], pos[0] / window_width)
                        self.assertAlmostEqual(out_value_attr.get()[1], pos[1] / window_height)
                    # Non-activated codepath.
                    else:
                        self.assertEqual(out_pressed_attr.get(), self.D)
                        self.assertEqual(out_released_attr.get(), self.D)
                        self.assertEqual(out_valuechanged_attr.get(), self.D)
                        self.assertFalse(out_ispressed_attr.get())
                        self.assertEqual(len(out_value_attr.get()), 2)
                        self.assertEqual(out_value_attr.get()[0], 0)
                        self.assertEqual(out_value_attr.get()[1], 0)

        # Test that the OnMouseInput node works correctly when its onlyPlayback input is disabled.
        in_onlyplayback_attr.set(False)
        for _ in _test_onmouseinput_node():
            yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.

        # Test that the OnMouseInput node works correctly when its onlyPlayback input is enabled.
        in_onlyplayback_attr.set(True)
        timeline = omni.timeline.get_timeline_interface()
        timeline.set_target_framerate(timeline.get_time_codes_per_seconds())
        timeline.play()
        for _ in _test_onmouseinput_node():
            yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
        timeline.stop()

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

        cube = ThreadsafetyTestUtils.add_to_threading_cache(
            test_instance_id, ogts.create_cube(stage, "Cube", (0.6, 0.4, 0.0))
        )
        attr_position = ThreadsafetyTestUtils.add_to_threading_cache(
            test_instance_id, cube.CreateAttribute("xformOp:translate", Sdf.ValueTypeNames.Float3, False)
        )
        attr_rotation = ThreadsafetyTestUtils.add_to_threading_cache(
            test_instance_id, cube.CreateAttribute("xformOp:rotateXYZ", Sdf.ValueTypeNames.Float3, False)
        )

        # 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, (on_object_change_node, flip_flop_node), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    (graph_path + "/OnObjectChange", "omni.graph.action.OnObjectChange"),
                    (graph_path + "/FlipFlop", "omni.graph.action.FlipFlop"),
                ],
                self.keys.SET_VALUES: [
                    (graph_path + "/OnObjectChange.inputs:onlyPlayback", False),
                    (graph_path + "/OnObjectChange.inputs:prim", attr_position.GetPrimPath()),
                    (graph_path + "/OnObjectChange.inputs:name", attr_position.GetName()),
                ],
                self.keys.CONNECT: (
                    graph_path + "/OnObjectChange.outputs:changed",
                    graph_path + "/FlipFlop.inputs:execIn",
                ),
            },
        )

        outa = og.Controller.attribute("outputs:a", flip_flop_node)

        # Check the current FlipFlop state, and that it isn't changing until we move the cube.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        ff_a_state_1 = og.Controller.get(outa)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(og.Controller.get(outa), ff_a_state_1)

        # Try changing a different (non-translate) attribute - should not trigger.
        ThreadsafetyTestUtils.single_evaluation_first_test_instance(
            test_instance_id, partial(attr_rotation.Set, Gf.Vec3f(180.0, 0.0, 0.0))
        )
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(og.Controller.get(outa), ff_a_state_1)

        # Now change the translate attribute - should trigger.
        ThreadsafetyTestUtils.single_evaluation_first_test_instance(
            test_instance_id, partial(attr_position.Set, Gf.Vec3f(1.0, 0.0, 0.0))
        )
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.

        ff_a_state_2 = og.Controller.get(outa)
        self.assertEqual(ff_a_state_2, 1 if not ff_a_state_1 else 0)
        property_path = og.Controller.get(og.Controller.attribute("outputs:propertyName", on_object_change_node))
        self.assertEqual(property_path, attr_position.GetName())

        # Look at the prim itself.
        og.Controller.edit(
            graph,
            {self.keys.SET_VALUES: (graph_path + "/OnObjectChange.inputs:name", "")},
        )
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.

        # Now change the rotate attribute - should trigger.
        ThreadsafetyTestUtils.single_evaluation_first_test_instance(
            test_instance_id, partial(attr_rotation.Set, Gf.Vec3f(245.0, 0.0, 0.0))
        )
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.

        ff_a_state_1 = og.Controller.get(outa)
        self.assertEqual(ff_a_state_1, 1 if not ff_a_state_2 else 0)
        property_path = og.Controller.get(og.Controller.attribute("outputs:propertyName", on_object_change_node))
        self.assertEqual(property_path, attr_rotation.GetName())

        # Now use the path input instead of the target. Do this for all OnObjectChange nodes!
        og.Controller.edit(
            graph,
            {
                self.keys.SET_VALUES: [
                    (graph_path + "/OnObjectChange.inputs:path", cube.GetPath().pathString),
                    (graph_path + "/OnObjectChange.inputs:prim", []),
                ]
            },
        )
        # Compute once to set up the inputs.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        # Change the rotate attribute - should trigger.
        ThreadsafetyTestUtils.single_evaluation_first_test_instance(
            test_instance_id, partial(attr_rotation.Set, Gf.Vec3f(0.0, 0.0, 0.0))
        )
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.

        ff_a_state_2 = og.Controller.get(outa)
        self.assertEqual(ff_a_state_2, 1 if not ff_a_state_1 else 0)
        property_path = og.Controller.get(og.Controller.attribute("outputs:propertyName", on_object_change_node))
        self.assertEqual(property_path, attr_rotation.GetName())

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

        timeline = omni.timeline.get_timeline_interface()

        # The stage's default stage is probably 60, set to 30 for this specific case
        stage = omni.usd.get_context().get_stage()
        fps = 30.0
        stage.SetTimeCodesPerSecond(fps)
        stage.SetStartTimeCode(1.0)
        stage.SetEndTimeCode(10.0)

        # 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_p_tick_node, _), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnPTick", "omni.graph.action.OnPlaybackTick"),
                    ("FlipFlop", "omni.graph.action.FlipFlop"),
                ],
                self.keys.CONNECT: [
                    ("OnPTick.outputs:tick", "FlipFlop.inputs:execIn"),
                ],
            },
        )

        # Obtain necessary attributes.
        tick_time_attr = on_p_tick_node.get_attribute("outputs:time")
        tick_frame_attr = on_p_tick_node.get_attribute("outputs:frame")

        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.

        # Check that the OnPlaybackTick node doesn't trigger when playback is not active.
        self.assertEqual(tick_time_attr.get(), 0)
        self.assertEqual(tick_frame_attr.get(), 0)

        # Check that the OnPlaybackTick node does trigger when playback is active, and the values are correct.
        timeline.play()
        yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertAlmostEqual(tick_time_attr.get(), timeline.get_current_time(), places=5)
        self.assertAlmostEqual(tick_frame_attr.get(), tick_time_attr.get() * fps, places=5)

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_ontick_node(self, test_instance_id: int = 0):
        """Test OnTick node"""
        timeline = omni.timeline.get_timeline_interface()

        # The stage's default stage is probably 60, set to 30 for this specific case
        stage = omni.usd.get_context().get_stage()
        fps = 30.0
        stage.SetTimeCodesPerSecond(fps)
        stage.SetStartTimeCode(1.0)
        stage.SetEndTimeCode(10.0)

        # 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, _), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("FlipFlop", "omni.graph.action.FlipFlop"),
                ],
                self.keys.CONNECT: [
                    ("OnTick.outputs:tick", "FlipFlop.inputs:execIn"),
                ],
                self.keys.SET_VALUES: ("OnTick.inputs:onlyPlayback", True),
            },
        )

        # Obtain necessary attributes.
        ontick_time_attr = on_tick_node.get_attribute("outputs:time")
        ontick_frame_attr = on_tick_node.get_attribute("outputs:frame")

        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.

        # Check that the OnTick node doesn't trigger when playback is not active.
        self.assertEqual(ontick_time_attr.get(), 0)
        self.assertEqual(ontick_frame_attr.get(), 0)

        # Check that the OnTick node does trigger when playback is active, and the values are correct.
        timeline.play()
        yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertAlmostEqual(ontick_time_attr.get(), timeline.get_current_time(), places=5)
        self.assertAlmostEqual(ontick_frame_attr.get(), ontick_time_attr.get() * fps, places=5)

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_onvariablechange_node(self, test_instance_id: int = 0):
        """Test OnVariableChange 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, (on_variable_change_node, counter_node), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnVariableChange", "omni.graph.action.OnVariableChange"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                self.keys.SET_VALUES: [
                    ("OnVariableChange.inputs:onlyPlayback", False),
                ],
                self.keys.CONNECT: [("OnVariableChange.outputs:changed", "Counter.inputs:execIn")],
            },
        )
        graph_context = graph.get_default_graph_context()

        # Create a new float variable on each graph, and set the OnVariableChange nodes to track this single variable.
        var_name = "myVar" + str(test_instance_id)
        my_var = og.Controller.create_variable(graph, var_name, "FLOAT")
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        og.Controller.set(og.Controller.attribute("inputs:variableName", on_variable_change_node), var_name)

        # Check that the execution output is still set to zero.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:changed", on_variable_change_node)), self.D)

        # Change the variable's value, check that the execution reads true now.
        my_var.set(graph_context, 5.2)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:changed", on_variable_change_node)), self.E)

        # Evaluate a second time to check that the execution gets set back to zero.
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:changed", on_variable_change_node)), self.D)

        # Make sure that the output execution remains off if we try setting the variable
        # to the same previous value.
        my_var.set(graph_context, 5.2)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:changed", on_variable_change_node)), self.D)

        # Check that setting the inputs:variableName to a nonexistant variable results in the output
        # execution pin remaining zero.
        og.Controller.set(og.Controller.attribute("inputs:variableName", on_variable_change_node), "nonexistant_name")
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:changed", on_variable_change_node)), self.D)

        # Set the onlyPlayback flag to true, run a similar set of tests to ensure that this setting
        # still works with the node. Note that a Counter node is used here to detect the variable change
        # because it's trickier to check against outputs:changed directly (harder to time everything with
        # all of the async stuff + when the timeline is running).

        og.Controller.set(og.Controller.attribute("inputs:variableName", on_variable_change_node), var_name)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        og.Controller.set(og.Controller.attribute("inputs:onlyPlayback", on_variable_change_node), True)
        timeline = ThreadsafetyTestUtils.add_to_threading_cache(
            test_instance_id, omni.timeline.get_timeline_interface()
        )
        ThreadsafetyTestUtils.single_evaluation_first_test_instance(
            test_instance_id, partial(timeline.set_target_framerate, timeline.get_time_codes_per_seconds())
        )
        ThreadsafetyTestUtils.single_evaluation_first_test_instance(test_instance_id, partial(timeline.play))
        my_var.set(graph_context, 4.6)
        yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:count", counter_node)), 2)
        my_var.set(graph_context, 4.6)
        yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:count", counter_node)), 2)
        yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:count", counter_node)), 2)
        my_var.set(graph_context, 4.6)
        yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:count", counter_node)), 2)
        og.Controller.set(og.Controller.attribute("inputs:variableName", on_variable_change_node), "nonexistant_name_2")
        yield ThreadsafetyTestUtils.EVALUATION_WAIT_FRAME  # Yielding to compute by waiting for the next app frame.
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:count", counter_node)), 2)
        ThreadsafetyTestUtils.single_evaluation_last_test_instance(test_instance_id, partial(timeline.stop))

    # ----------------------------------------------------------------------
    async def _run_test_variable_of_type(self, graph_id, var_type: og.Type, initial_value, changed_value):
        """Helper method to run a variable test with different variable types"""
        graph_path = self.TEST_GRAPH_PATH + graph_id
        og.Controller.create_graph({"graph_path": graph_path, "evaluator_name": "execution"})
        (graph, (_, counter_node), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnVariableChange", "omni.graph.action.OnVariableChange"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                self.keys.SET_VALUES: [
                    ("OnVariableChange.inputs:onlyPlayback", False),
                    ("OnVariableChange.inputs:variableName", "var"),
                ],
                self.keys.CONNECT: [("OnVariableChange.outputs:changed", "Counter.inputs:execIn")],
                self.keys.CREATE_VARIABLES: [
                    ("var", var_type),
                ],
            },
        )

        # test that nothing triggered on the initial run
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:count", counter_node)), 0)
        var = graph.find_variable("var")

        # set the initial value, expected a trigger
        var.set(graph.get_context(), initial_value)
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:count", counter_node)), 1)
        ogts.verify_values(var.get(graph.get_context()), initial_value, "expected initial value to be set")

        # change the value of the variable and makes sure the OnVariableChange triggers
        var.set(graph.get_context(), changed_value)
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(og.Controller.get(og.Controller.attribute("outputs:count", counter_node)), 2)
        ogts.verify_values(var.get(graph.get_context()), changed_value, "expeected changed value to be set")

    # ----------------------------------------------------------------------
    async def test_onvariablechanged_by_type(self):
        """Tests OnVariableChange for different types of variables, including strings and arrays"""
        await self._run_test_variable_of_type("_int", og.Type(og.BaseDataType.INT), 1, 2)
        await self._run_test_variable_of_type("_int_array", og.Type(og.BaseDataType.INT, 1, 1), [1, 2, 3], [3, 4, 5])
        await self._run_test_variable_of_type(
            "_int_array_size", og.Type(og.BaseDataType.INT, 1, 1), [1, 2, 3, 4, 5], [1, 2, 3]
        )
        await self._run_test_variable_of_type("_tuple", og.Type(og.BaseDataType.INT, 3, 0), (1, 2, 3), (3, 4, 5))
        await self._run_test_variable_of_type(
            "_string", og.Type(og.BaseDataType.UCHAR, 1, 1, og.AttributeRole.TEXT), "init", "changed"
        )
        await self._run_test_variable_of_type(
            "_tuple_array", og.Type(og.BaseDataType.INT, 3, 1), [(1, 2, 3), (4, 5, 6)], [(3, 4, 5)]
        )

    # ----------------------------------------------------------------------
    async def test_onvariablechange_node_instances(self):
        """Tests that OnVariableChange node works with instancing"""
        controller = og.Controller()
        keys = og.Controller.Keys
        int_range = range(1, 100)

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

        for i in int_range:
            prim_name = f"/World/Prim_{i}"
            prim = stage.DefinePrim(prim_name)
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim_name, self.TEST_GRAPH_PATH)
            prim.CreateAttribute("graph:variable:int_var", Sdf.ValueTypeNames.Int).Set(0)
            prim.CreateAttribute("graph_output", Sdf.ValueTypeNames.Int).Set(i)

        (_, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnVariableChange", "omni.graph.action.OnVariableChange"),
                    ("GraphTarget", "omni.graph.nodes.GraphTarget"),
                    ("ReadPrimAttr", "omni.graph.nodes.ReadPrimAttribute"),
                    ("WritePrimAttr", "omni.graph.nodes.WritePrimAttribute"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("ConstantInt", "omni.graph.nodes.ConstantInt"),
                ],
                keys.CONNECT: [
                    ("OnVariableChange.outputs:changed", "WritePrimAttr.inputs:execIn"),
                    ("GraphTarget.outputs:primPath", "WritePrimAttr.inputs:primPath"),
                    ("GraphTarget.outputs:primPath", "ReadPrimAttr.inputs:primPath"),
                    ("ReadPrimAttr.outputs:value", "Add.inputs:a"),
                    ("ConstantInt.inputs:value", "Add.inputs:b"),
                    ("Add.outputs:sum", "WritePrimAttr.inputs:value"),
                ],
                keys.CREATE_VARIABLES: [("int_var", og.Type(og.BaseDataType.INT))],
                keys.SET_VALUES: [
                    ("OnVariableChange.inputs:onlyPlayback", False),
                    ("OnVariableChange.inputs:variableName", "int_var"),
                    ("WritePrimAttr.inputs:usePath", True),
                    ("WritePrimAttr.inputs:name", "graph_output"),
                    ("ReadPrimAttr.inputs:usePath", True),
                    ("ReadPrimAttr.inputs:name", "graph_output"),
                    ("ConstantInt.inputs:value", 1),
                ],
            },
        )
        await omni.kit.app.get_app().next_update_async()

        # Setting value of int_var from 0 to 1.
        for i in int_range:
            prim_path = f"/World/Prim_{i}"
            prev_val = stage.GetPrimAtPath(prim_path).GetAttribute("graph_output").Get()
            stage.GetPrimAtPath(prim_path).GetAttribute("graph:variable:int_var").Set(1)
            await omni.kit.app.get_app().next_update_async()

            val = stage.GetPrimAtPath(prim_path).GetAttribute("graph_output").Get()
            self.assertEqual(val, prev_val + 1, msg=f"{prim_path}")

        # Setting value of int_var from 1 to 1.
        for i in int_range:
            prim_path = f"/World/Prim_{i}"
            prev_val = stage.GetPrimAtPath(prim_path).GetAttribute("graph_output").Get()
            stage.GetPrimAtPath(prim_path).GetAttribute("graph:variable:int_var").Set(1)
            await omni.kit.app.get_app().next_update_async()

            val = stage.GetPrimAtPath(prim_path).GetAttribute("graph_output").Get()
            self.assertEqual(val, prev_val + 1, msg=f"{prim_path}")

    # ----------------------------------------------------------------------
    # The SendCustomEvent node is tested in tandem with the OnCustomEvent node above; it is also used
    # extensively in other tests, so we skip adding another dedicated test here.

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_sequence_node(self, test_instance_id: int = 0):
        """Test Sequence 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"})
        (_, (_, sequence_node, counter0_node, counter1_node), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Sequence", "omni.graph.action.Sequence"),
                    ("Counter0", "omni.graph.action.Counter"),
                    ("Counter1", "omni.graph.action.Counter"),
                ],
                self.keys.SET_VALUES: [("OnTick.inputs:onlyPlayback", False)],
                self.keys.CONNECT: [
                    ("OnTick.outputs:tick", "Sequence.inputs:execIn"),
                    ("Sequence.outputs:a", "Counter0.inputs:execIn"),
                    ("Sequence.outputs:b", "Counter1.inputs:execIn"),
                ],
            },
        )

        # Obtain necessary attributes.
        in_exec_attr_0 = counter0_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_a_attr = sequence_node.get_attribute("outputs:a")
        out_b_attr = sequence_node.get_attribute("outputs:b")

        # Check that the Sequence 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 Sequence
        # nodes remain enabled (regardless of the fact that they're not connected downstream).
        # Similar test to the Multisequence node.
        self.assertEqual(out_cnt_attr_0.get(), 0)
        self.assertEqual(out_cnt_attr_1.get(), 0)
        self.assertEqual(out_a_attr.get(), self.D)
        self.assertEqual(out_b_attr.get(), self.D)
        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_a_attr.get(), self.D)
        self.assertEqual(out_b_attr.get(), self.E)
        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_a_attr.get(), self.D)
        self.assertEqual(out_b_attr.get(), self.E)

        # Connect output pin b to the Counter1 node. This is because both pins a and b share
        # the same terminal input attribute, so both must have the same value. Since
        # b is the last attribute and must be set to 1 (due to logic inside this node's
        # compute method), output a must also take on this value.
        og.Controller.connect(out_b_attr, in_exec_attr_0)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_cnt_attr_0.get(), 4)
        self.assertEqual(out_cnt_attr_1.get(), 3)
        self.assertEqual(out_a_attr.get(), self.E)
        self.assertEqual(out_b_attr.get(), self.E)

    # ----------------------------------------------------------------------
    async def test_setprimactive_node(self):
        await self._setprimactive_node(False)

    # ----------------------------------------------------------------------
    async def test_setprimactive_node_with_target(self):
        await self._setprimactive_node(True)

    # ----------------------------------------------------------------------
    async def _setprimactive_node(self, use_target_inputs):
        """
        Test SetPrimActive node. If use_target_inputs is true, will use the prim target inputs rather than prim path
        """
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        (_, (_, set_prim_active_node), _, _,) = og.Controller.edit(
            self.TEST_GRAPH_PATH,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("SetPrimActive", "omni.graph.action.SetPrimActive"),
                ],
                self.keys.CREATE_PRIMS: [
                    ("/TestPrim", {}),
                ],
                self.keys.CONNECT: [("OnTick.outputs:tick", "SetPrimActive.inputs:execIn")],
                self.keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("SetPrimActive.inputs:prim", "/TestPrim"),
                    ("SetPrimActive.inputs:active", True),
                ],
            },
        )
        await og.Controller.evaluate()

        if use_target_inputs:
            rel = stage.GetRelationshipAtPath(f"{self.TEST_GRAPH_PATH}/SetPrimActive.inputs:primTarget")
            omni.kit.commands.execute("AddRelationshipTarget", relationship=rel, target="/TestPrim")
            await og.Controller.evaluate()

        # Check that the prim we're pointing to remains active when we trigger
        # graph evaluation.
        prim = stage.GetPrimAtPath("/TestPrim")
        self.assertTrue(prim.IsActive())
        await og.Controller.evaluate()
        self.assertTrue(prim.IsActive())

        # Now set the prim to be inactive.
        og.Controller.set(og.Controller.attribute("inputs:active", set_prim_active_node), False)
        await og.Controller.evaluate()
        self.assertFalse(prim.IsActive())
        await og.Controller.evaluate()
        self.assertFalse(prim.IsActive())

        # Again set to be active, make sure that the prim gets reactivated.
        og.Controller.set(og.Controller.attribute("inputs:active", set_prim_active_node), True)
        await og.Controller.evaluate()
        self.assertTrue(prim.IsActive())

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_switchtoken_node(self, test_instance_id: int = 0):
        """Test SwitchToken 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"})
        (_, (_, switch_node), _, _,) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Switch", "omni.graph.action.SwitchToken"),
                ],
                self.keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                ],
                self.keys.CONNECT: [
                    ("OnTick.outputs:tick", "Switch.inputs:execIn"),
                ],
            },
        )

        # Test that an execution connection is correctly triggering the downstream nodes once per evaluation.
        def add_input(index: int):
            og.Controller.create_attribute(
                switch_node,
                f"inputs:branch{index:02}",
                "token",
                og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
            )
            og.Controller.create_attribute(
                switch_node,
                f"outputs:output{index:02}",
                "execution",
                og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT,
            )
            og.Controller.set(og.Controller.attribute(f"inputs:branch{index:02}", switch_node), f"{index:02}")

        for index in range(5):
            add_input(index)

        for index in range(5):
            og.Controller.set(og.Controller.attribute("inputs:value", switch_node), f"{index:02}")
            yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
            for index2 in range(5):
                expected_val = (
                    og.ExecutionAttributeState.DISABLED if index2 != index else og.ExecutionAttributeState.ENABLED
                )
                self.assertEqual(
                    og.Controller.get(og.Controller.attribute(f"outputs:output{index2:02}", switch_node)),
                    expected_val,
                )

    # ----------------------------------------------------------------------
    @ThreadsafetyTestUtils.make_threading_test
    def test_syncgate_node(self, test_instance_id: int = 0):
        """Test SyncGate 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_impulse0_node, on_impulse1_node, on_impulse2_node, on_impulse3_node, sync_gate_node, counter_node),
            _,
            _,
        ) = og.Controller.edit(
            graph_path,
            {
                self.keys.CREATE_NODES: [
                    ("OnImpulse0", "omni.graph.action.OnImpulseEvent"),
                    ("OnImpulse1", "omni.graph.action.OnImpulseEvent"),
                    ("OnImpulse2", "omni.graph.action.OnImpulseEvent"),
                    ("OnImpulse3", "omni.graph.action.OnImpulseEvent"),
                    ("SyncGate", "omni.graph.action.SyncGate"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                self.keys.CONNECT: [
                    ("OnImpulse0.outputs:execOut", "SyncGate.inputs:execIn"),
                    ("OnImpulse1.outputs:execOut", "SyncGate.inputs:execIn"),
                    ("OnImpulse2.outputs:execOut", "SyncGate.inputs:execIn"),
                    ("OnImpulse3.outputs:execOut", "SyncGate.inputs:execIn"),
                    ("SyncGate.outputs:execOut", "Counter.inputs:execIn"),
                ],
                self.keys.SET_VALUES: [
                    ("OnImpulse0.inputs:onlyPlayback", False),
                    ("OnImpulse1.inputs:onlyPlayback", False),
                    ("OnImpulse2.inputs:onlyPlayback", False),
                    ("OnImpulse3.inputs:onlyPlayback", False),
                    ("SyncGate.inputs:syncValue", 5),
                ],
            },
        )

        # Obtain necessary attributes.
        state_enable_impulse_attr_0 = on_impulse0_node.get_attribute("state:enableImpulse")
        state_enable_impulse_attr_1 = on_impulse1_node.get_attribute("state:enableImpulse")
        state_enable_impulse_attr_2 = on_impulse2_node.get_attribute("state:enableImpulse")
        state_enable_impulse_attr_3 = on_impulse3_node.get_attribute("state:enableImpulse")
        in_syncvalue_attr = sync_gate_node.get_attribute("inputs:syncValue")
        out_syncvalue_attr = sync_gate_node.get_attribute("outputs:syncValue")
        out_cnt_attr = counter_node.get_attribute("outputs:count")

        # First check that the SyncGate node only gets unblocked when it detects 4+ input
        # enabled executions, at which point the Counter node will begin to be incremented.
        # Also check that after the 1st graph evaluation the SyncGate's output syncValue gets
        # set to its input syncValue.
        state_enable_impulse_attr_0.set(True)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_syncvalue_attr.get(), 5)
        self.assertEqual(out_cnt_attr.get(), 0)

        state_enable_impulse_attr_1.set(True)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_syncvalue_attr.get(), 5)
        self.assertEqual(out_cnt_attr.get(), 0)

        state_enable_impulse_attr_2.set(True)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_syncvalue_attr.get(), 5)
        self.assertEqual(out_cnt_attr.get(), 0)

        state_enable_impulse_attr_3.set(True)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_syncvalue_attr.get(), 5)
        self.assertEqual(out_cnt_attr.get(), 1)

        state_enable_impulse_attr_0.set(True)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_syncvalue_attr.get(), 5)
        self.assertEqual(out_cnt_attr.get(), 2)

        # Now reset the SyncGate node's syncValue. Check that we once again need 4+ input executions
        # with the "enabled" status flowing into the SyncGate in order to unlock it. Also note
        # that we don't technically need each separate OnImpulse node to trigger to unlock the gate;
        # we could have a single OnImpulse node (per graph) send 4 separate impulses to open each gate
        # in each graph, since that would reach the threshold number of accumulated execIn states for
        # this setup (equal to 4 since we have 4 nodes connected to the SyncGate's inputs:execIn pin).
        in_syncvalue_attr.set(9)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        for i in range(0, 5):
            state_enable_impulse_attr_0.set(True)
            yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
            self.assertEqual(out_syncvalue_attr.get(), 9)
            if i < 3:
                self.assertEqual(out_cnt_attr.get(), 2)
            else:
                self.assertEqual(out_cnt_attr.get(), 2 + (i - 2))

        # Now reset the syncValue again per node. This time only send 2 input impulses, then reset the syncValue
        # immediately after. Check that we'll again need to send 4 input impulses to open the corresponding gate
        # (which was essentially "relocked" with the syncValue change).
        in_syncvalue_attr.set(21)
        state_enable_impulse_attr_0.set(True)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_syncvalue_attr.get(), 21)
        self.assertEqual(out_cnt_attr.get(), 4)

        state_enable_impulse_attr_2.set(True)
        yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
        self.assertEqual(out_syncvalue_attr.get(), 21)
        self.assertEqual(out_cnt_attr.get(), 4)

        in_syncvalue_attr.set(25)
        for i in range(0, 5):
            state_enable_impulse_attr_2.set(True)
            yield  # Yielding to wait for compute to happen across all graph instances before continuing the test.
            self.assertEqual(out_syncvalue_attr.get(), 25)
            if i < 3:
                self.assertEqual(out_cnt_attr.get(), 4)
            else:
                self.assertEqual(out_cnt_attr.get(), 4 + (i - 2))

    # ----------------------------------------------------------------------
    async def test_countdown_node(self):
        """Test Countdown node"""
        duration = 5
        period = 3

        og.Controller.create_graph({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})
        (_, (_, _, counter0_node, counter1_node), _, _,) = og.Controller.edit(
            self.TEST_GRAPH_PATH,
            {
                self.keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Countdown", "omni.graph.action.Countdown"),
                    ("Counter0", "omni.graph.action.Counter"),
                    ("Counter1", "omni.graph.action.Counter"),
                ],
                self.keys.CONNECT: [
                    ("OnTick.outputs:tick", "Countdown.inputs:execIn"),
                    ("Countdown.outputs:finished", "Counter0.inputs:execIn"),
                    ("Countdown.outputs:tick", "Counter1.inputs:execIn"),
                ],
                # Set the Countdown inputs so that the Counter0 node only gets incremented
                # once 4 update ticks have passed since the graph creation, and so that
                # the Counter1 node gets incremented every 2 ticks.
                self.keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("Countdown.inputs:duration", duration),
                    ("Countdown.inputs:period", period),
                ],
            },
        )

        # Obtain necessary attributes.
        out_cnt_attr_0 = counter0_node.get_attribute("outputs:count")
        out_cnt_attr_1 = counter1_node.get_attribute("outputs:count")

        # Test against a predetermined set of values to make sure that
        # this node doesn't break in the future.
        cnt0 = [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2]
        cnt1 = [0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3]
        for i in range(0, 20):
            await og.Controller.evaluate()
            self.assertEqual(og.Controller.get(out_cnt_attr_0), cnt0[i])
            self.assertEqual(og.Controller.get(out_cnt_attr_1), cnt1[i])
