"""Basic tests of the execution framework with OG"""

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


# ======================================================================
class SimpleGraphScene:
    def __init__(self):
        self.graph = None
        self.nodes = None


# ======================================================================
def create_simple_graph():
    """Create a scene containing a PushGraph with a TestExecutionTask."""

    scene = SimpleGraphScene()
    (scene.graph, scene.nodes, _, _) = og.Controller.edit(
        "/World/PushGraph",
        {
            og.Controller.Keys.CREATE_NODES: [
                ("TestNode1", "omni.graph.test.TestExecutionTask"),
                ("TestNode2", "omni.graph.test.TestExecutionTask"),
                ("TestNode3", "omni.graph.test.TestExecutionTask"),
                ("TestNode4", "omni.graph.test.TestExecutionTask"),
            ],
        },
    )

    return scene


# ======================================================================
def create_concurrency_graph():
    """Create a scene containing a PushGraph with a TestConcurrency."""

    scene = SimpleGraphScene()
    (scene.graph, scene.nodes, _, _) = og.Controller.edit(
        "/World/PushGraph",
        {
            og.Controller.Keys.CREATE_NODES: [
                ("TestNode1", "omni.graph.test.TestConcurrency"),
                ("TestNode2", "omni.graph.test.TestConcurrency"),
                ("TestNode3", "omni.graph.test.TestConcurrency"),
                ("TestNode4", "omni.graph.test.TestConcurrency"),
            ],
            og.Controller.Keys.SET_VALUES: [
                ("TestNode1.inputs:expected", 4),
                ("TestNode2.inputs:expected", 4),
                ("TestNode3.inputs:expected", 4),
                ("TestNode4.inputs:expected", 4),
                ("TestNode1.inputs:timeOut", 100),
                ("TestNode2.inputs:timeOut", 100),
                ("TestNode3.inputs:timeOut", 100),
                ("TestNode4.inputs:timeOut", 100),
            ],
        },
    )

    return scene


# ======================================================================
def create_isolate_graph():
    """Create a scene containing a PushGraph with TestIsolate, TestSerial, and TestConcurrency nodes."""

    scene = SimpleGraphScene()
    (scene.graph, scene.nodes, _, _) = og.Controller.edit(
        "/World/PushGraph",
        {
            og.Controller.Keys.CREATE_NODES: [
                ("TC0", "omni.graph.test.TestConcurrency"),
                ("TC1", "omni.graph.test.TestConcurrency"),
                ("TC2", "omni.graph.test.TestConcurrency"),
                ("TC3", "omni.graph.test.TestConcurrency"),
                ("TC4", "omni.graph.test.TestConcurrency"),
                ("TI0", "omni.graph.test.TestIsolate"),
                ("TI1", "omni.graph.test.TestIsolate"),
                ("TI2", "omni.graph.test.TestIsolate"),
                ("TS0", "omni.graph.test.TestSerial"),
                ("TS1", "omni.graph.test.TestSerial"),
                ("TS2", "omni.graph.test.TestSerial"),
                ("TS3", "omni.graph.test.TestSerial"),
                ("TS4", "omni.graph.test.TestSerial"),
            ],
            og.Controller.Keys.SET_VALUES: [
                ("TC0.inputs:expected", 5),
                ("TC1.inputs:expected", 5),
                ("TC2.inputs:expected", 5),
                ("TC3.inputs:expected", 5),
                ("TC4.inputs:expected", 5),
                ("TC0.inputs:timeOut", 100),
                ("TC1.inputs:timeOut", 100),
                ("TC2.inputs:timeOut", 100),
                ("TC3.inputs:timeOut", 100),
                ("TC4.inputs:timeOut", 100),
            ],
        },
    )

    return scene


# ======================================================================
def create_cyclic_graph(evaluator_name: str):
    """Create a scene containing a graph with cycles."""

    # 4 cycles exist, which are made up of:
    # 1. Nodes 8 and 9.
    # 2. Nodes 11 and 12.
    # 3. Nodes 11, 13, and 15.
    # 4. Nodes 11, 13, 14, and 15.
    # Note that we also use a mix of serial- and parallel-scheduled
    # nodes for this test graph, just for fun :)
    scene = SimpleGraphScene()
    (scene.graph, scene.nodes, _, _) = og.Controller.edit(
        {"graph_path": "/World/TestGraph", "evaluator_name": evaluator_name},
        {
            og.Controller.Keys.CREATE_NODES: [
                ("Node0", "omni.graph.test.TestCyclesSerial"),
                ("Node1", "omni.graph.test.TestCyclesParallel"),
                ("Node2", "omni.graph.test.TestCyclesParallel"),
                ("Node3", "omni.graph.test.TestCyclesSerial"),
                ("Node4", "omni.graph.test.TestCyclesSerial"),
                ("Node5", "omni.graph.test.TestCyclesParallel"),
                ("Node6", "omni.graph.test.TestCyclesSerial"),
                ("Node7", "omni.graph.test.TestCyclesParallel"),
                ("Node8", "omni.graph.test.TestCyclesParallel"),
                ("Node9", "omni.graph.test.TestCyclesParallel"),
                ("Node10", "omni.graph.test.TestCyclesParallel"),
                ("Node11", "omni.graph.test.TestCyclesParallel"),
                ("Node12", "omni.graph.test.TestCyclesParallel"),
                ("Node13", "omni.graph.test.TestCyclesParallel"),
                ("Node14", "omni.graph.test.TestCyclesSerial"),
                ("Node15", "omni.graph.test.TestCyclesSerial"),
                ("Node16", "omni.graph.test.TestCyclesSerial"),
            ],
            og.Controller.Keys.CONNECT: [
                ("Node0.outputs:a", "Node4.inputs:a"),
                ("Node4.outputs:a", "Node5.inputs:a"),
                ("Node5.outputs:a", "Node2.inputs:a"),
                ("Node5.outputs:b", "Node6.inputs:a"),
                ("Node2.outputs:a", "Node3.inputs:a"),
                ("Node6.outputs:a", "Node3.inputs:b"),
                ("Node1.outputs:a", "Node2.inputs:b"),
                ("Node0.outputs:b", "Node8.inputs:a"),
                ("Node8.outputs:a", "Node9.inputs:a"),
                ("Node9.outputs:a", "Node8.inputs:b"),
                ("Node9.outputs:b", "Node7.inputs:a"),
                ("Node3.outputs:a", "Node7.inputs:b"),
                ("Node0.outputs:c", "Node10.inputs:a"),
                ("Node10.outputs:a", "Node11.inputs:a"),
                ("Node11.outputs:a", "Node12.inputs:a"),
                ("Node12.outputs:a", "Node11.inputs:b"),
                ("Node11.outputs:b", "Node16.inputs:a"),
                ("Node11.outputs:c", "Node13.inputs:a"),
                ("Node13.outputs:a", "Node14.inputs:a"),
                ("Node14.outputs:a", "Node15.inputs:a"),
                ("Node13.outputs:b", "Node15.inputs:b"),
                ("Node15.outputs:a", "Node11.inputs:c"),
                ("Node0.outputs:d", "Node0.inputs:d"),
                ("Node5.outputs:c", "Node5.inputs:c"),
                ("Node10.outputs:b", "Node10.inputs:b"),
                ("Node14.outputs:b", "Node14.inputs:b"),
            ],
        },
    )

    return scene


# ======================================================================
def create_single_node_cyclic_graph(evaluator_name: str):
    """Create a scene containing a graph with single-node cycles"""
    scene = SimpleGraphScene()
    (scene.graph, scene.nodes, _, _) = og.Controller.edit(
        {"graph_path": "/World/TestGraph", "evaluator_name": evaluator_name},
        {
            og.Controller.Keys.CREATE_NODES: [
                ("Node0", "omni.graph.test.TestCyclesSerial"),
                ("Node1", "omni.graph.test.TestCyclesParallel"),
            ],
            og.Controller.Keys.CONNECT: [
                ("Node0.outputs:a", "Node0.inputs:a"),
                ("Node0.outputs:a", "Node1.inputs:a"),
                ("Node1.outputs:b", "Node1.inputs:b"),
            ],
        },
    )

    return scene


# ======================================================================
class TestExecutionFrameworkSanity(ogts.OmniGraphTestCase):
    """Execution Framework Unit Tests"""

    async def test_execution_task(self):
        """Validate that execution framework is executing nodes in the graph."""
        scene = create_simple_graph()

        # Verify starting condition.
        for node in scene.nodes:
            self.assertFalse(og.Controller(("outputs:result", node)).get())

        # Execute the update once and we want explicitly to execute entire app update pipeline.
        await omni.kit.app.get_app().next_update_async()

        # Verify all nodes were executed with EF.
        for node in scene.nodes:
            self.assertTrue(og.Controller(("outputs:result", node)).get())

    async def test_parallel_execution(self):
        """Validate that nodes are executed concurrently."""
        scene = create_concurrency_graph()

        def _check_all_reached_concurrency(scene):
            result = True
            for node in scene.nodes:
                result = result and og.Controller(("outputs:result", node)).get()

            return result

        # Verify starting condition.
        self.assertFalse(_check_all_reached_concurrency(scene))

        # First executions are forced to be run in isolation.
        await omni.kit.app.get_app().next_update_async()
        self.assertFalse(_check_all_reached_concurrency(scene))

        # Now we should be getting parallel execution.
        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(_check_all_reached_concurrency(scene))

        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(_check_all_reached_concurrency(scene))

        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(_check_all_reached_concurrency(scene))

    async def test_on_demand_execution_task(self):
        """Validate that execution framework is executing nodes in the graph."""
        scene = create_simple_graph()

        # Verify starting condition.
        for node in scene.nodes:
            self.assertFalse(og.Controller(("outputs:result", node)).get())

        # Execute given graph explicitly/on-demand.
        await og.Controller.evaluate(scene.graph)

        # Verify all nodes were executed with EF.
        for node in scene.nodes:
            self.assertTrue(og.Controller(("outputs:result", node)).get())

    async def test_isolate_execution(self):
        """Check that nodes with the "usd-write" scheduling hint get executed in isolation."""
        scene = create_isolate_graph()

        def _check_all_successful(scene):
            result = True
            for node in scene.nodes:
                result = result and og.Controller(("outputs:result", node)).get()

            return result

        # Verify starting condition.
        self.assertFalse(_check_all_successful(scene))

        # First executions are forced to be run in isolation.
        await omni.kit.app.get_app().next_update_async()

        # In subsequent executions each node gets executed per
        # its scheduling hint. Should not get any conflicts.
        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(_check_all_successful(scene))

        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(_check_all_successful(scene))

        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(_check_all_successful(scene))

    async def test_cycling_push_graph(self):
        """Check that cycles in push graphs correctly get ignored during execution."""
        scene = create_cyclic_graph("push")

        # Set of node indices that will evaluate.
        node_indices_that_evaluate = {0, 1, 2, 3, 4, 5, 6, 10}

        # Get all state counters.
        state_cnt_attrs = [node.get_attribute("state:count") for node in scene.nodes]

        # For each tick, only the nodes in the index list should get evaluated;
        # all other nodes either lie within a cycle or downstream of a cycle, and
        # thus will be ignored.
        for i in range(4):
            await omni.kit.app.get_app().next_update_async()
            for j, state_cnt_attr in enumerate(state_cnt_attrs):
                self.assertEqual(state_cnt_attr.get(), i + 1 if j in node_indices_that_evaluate else 0)

    async def test_cycling_pull_graph(self):
        """Check that cycles in pull graphs correctly get ignored during execution."""
        scene = create_cyclic_graph("dirty_push")

        # Set of node indices that will evaluate when node 0 is dirtied.
        node_indices_that_evaluate = {0, 1, 2, 3, 4, 5, 6, 10}

        # Get all state counters.
        state_cnt_attrs = [node.get_attribute("state:count") for node in scene.nodes]

        # Trigger a graph evaluation by adding/removing a connection between nodes 0 and 4.
        # Only the nodes in the index list should get evaluated; all other nodes either lie
        # within a cycle or downstream of a cycle, and thus will be ignored.
        for i in range(4):
            if i % 2 == 0:
                og.Controller.connect(
                    scene.nodes[0].get_attribute("outputs:d"), scene.nodes[4].get_attribute("inputs:d")
                )
            else:
                og.Controller.disconnect(
                    scene.nodes[0].get_attribute("outputs:d"), scene.nodes[4].get_attribute("inputs:d")
                )
            await omni.kit.app.get_app().next_update_async()
            for j, state_cnt_attr in enumerate(state_cnt_attrs):
                self.assertEqual(state_cnt_attr.get(), i + 1 if j in node_indices_that_evaluate else 0)

    async def test_push_graph_with_single_node_cycles(self):
        """Check that single-node cycles in push graphs do not block execution."""
        scene = create_single_node_cyclic_graph("push")

        # Get all state counters.
        state_cnt_attrs = [node.get_attribute("state:count") for node in scene.nodes]

        # Check that all nodes get executed each tick.
        for i in range(4):
            await omni.kit.app.get_app().next_update_async()
            for state_cnt_attr in state_cnt_attrs:
                self.assertEqual(state_cnt_attr.get(), i + 1)

    async def test_pull_graph_with_single_node_cycles(self):
        """Check that single-node cycles in pull graphs do not block execution."""
        scene = create_single_node_cyclic_graph("dirty_push")

        # Get all state counters.
        state_cnt_attrs = [node.get_attribute("state:count") for node in scene.nodes]

        # Trigger a graph evaluation by adding/removing an extra connection between nodes
        # 0 and 1. All nodes should execute.
        for i in range(4):
            if i % 2 == 0:
                og.Controller.connect(
                    scene.nodes[0].get_attribute("outputs:c"), scene.nodes[1].get_attribute("inputs:c")
                )
            else:
                og.Controller.disconnect(
                    scene.nodes[0].get_attribute("outputs:c"), scene.nodes[1].get_attribute("inputs:c")
                )
            await omni.kit.app.get_app().next_update_async()
            for state_cnt_attr in state_cnt_attrs:
                self.assertEqual(state_cnt_attr.get(), i + 1)

    async def test_partitioning_with_instancing(self):
        """
        Check that injection of either a NodeDefLambda or
        NodeGraphDef into the OG execution graph backend
        continues to work with instancing.
        """

        # Run the test against both test partition passes.
        partition_pass_settings = [
            "/app/omni.graph.test/enableTestNodeDefLambdaPass",
            "/app/omni.graph.test/enableTestNodeGraphDefPass",
        ]
        for si, pps in enumerate(partition_pass_settings):
            with og.Settings.temporary(pps, True):
                # Run the test over push and dirty_push evaluator types. Once partitioning
                # is supported by the Action Graph executor, we should be able to also
                # test the "execution" evaluator here as well.
                evaluator_names = [
                    "push",
                    "dirty_push",
                    # "execution",
                ]
                for evaluator_name in evaluator_names:
                    # Create our graph, which will contain a ReadVariable and
                    # WriteVariable node.
                    graph_path = "/World/TestGraph"
                    (graph, (read_variable_node, write_variable_node), _, _) = og.Controller.edit(
                        {"graph_path": graph_path, "evaluator_name": evaluator_name},
                        {
                            og.Controller.Keys.CREATE_NODES: [
                                ("ReadVariable", "omni.graph.core.ReadVariable"),
                                ("WriteVariable", "omni.graph.core.WriteVariable"),
                            ],
                            og.Controller.Keys.CREATE_VARIABLES: [
                                ("in_int_var", og.Type(og.BaseDataType.INT)),
                                ("out_int_var", og.Type(og.BaseDataType.INT)),
                            ],
                            og.Controller.Keys.CONNECT: [
                                ("ReadVariable.outputs:value", "WriteVariable.inputs:value"),
                            ],
                            og.Controller.Keys.SET_VALUES: [
                                ("ReadVariable.inputs:variableName", "in_int_var"),
                                ("WriteVariable.inputs:variableName", "out_int_var"),
                            ],
                        },
                    )

                    # Create the custom counter attribute on both nodes so
                    # that they get picked up by the partition pass.
                    og.Controller.create_attribute(
                        read_variable_node,
                        "outputs:test_counter_attribute",
                        og.Type(og.BaseDataType.INT, 1, 0, og.AttributeRole.NONE),
                        og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT,
                    )
                    og.Controller.create_attribute(
                        write_variable_node,
                        "outputs:test_counter_attribute",
                        og.Type(og.BaseDataType.INT, 1, 0, og.AttributeRole.NONE),
                        og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT,
                    )

                    # Get some attributes/variables that we'll need for the test.
                    in_var = graph.find_variable("in_int_var")
                    rn_out_value = read_variable_node.get_attribute("outputs:value")
                    wn_in_value = write_variable_node.get_attribute("inputs:value")
                    wn_out_value = write_variable_node.get_attribute("outputs:value")
                    out_var = graph.find_variable("out_int_var")
                    rn_cnt = read_variable_node.get_attribute("outputs:test_counter_attribute")
                    wn_cnt = write_variable_node.get_attribute("outputs:test_counter_attribute")

                    # Create 3 graph instances, each with unique "in_int_var" values
                    # and "out_int_var" set to zero.
                    num_prims = 3
                    prim_ids = [None] * num_prims
                    prim_paths = [None] * num_prims
                    prims = [None] * num_prims
                    stage = omni.usd.get_context().get_stage()
                    graph_context = graph.get_default_graph_context()
                    for i in range(0, num_prims):
                        prim_ids[i] = i
                        prim_paths[i] = f"/World/Prim_{i}"
                        prims[i] = stage.DefinePrim(prim_paths[i])
                        OmniGraphSchemaTools.applyOmniGraphAPI(stage, prim_paths[i], graph_path)
                        prims[i].CreateAttribute("graph:variable:in_int_var", Sdf.ValueTypeNames.Int).Set(i)
                        prims[i].CreateAttribute("graph:variable:out_int_var", Sdf.ValueTypeNames.Int).Set(0)

                    # When we evaluate the graph we should see that (a) the counter attribute on the
                    # authoring node gets ticked (which is handled directly by the partition passes,
                    # thus letting us know that they actually went through successfully/weren't just
                    # no-ops), and (b) the out_int_var variable gets successfully written out to with
                    # the value listed on the in_int_var variable.
                    await omni.kit.app.get_app().next_update_async()
                    await omni.kit.app.get_app().next_update_async()

                    for i in range(0, num_prims):
                        expected_value = i
                        self.assertEqual(in_var.get(graph_context, instance_path=prim_paths[i]), expected_value)
                        # Note that the instanced nodes (and thus the attributes we're looking up)
                        # need not be ordered according to the iteration index, e.g., instance with
                        # index 0 may correspond to the prim at the i == 2 index in the prim_paths
                        # list This is why we perform the subsequent 3 checks as shown below.
                        self.assertIn(rn_out_value.get(instance=i), prim_ids)
                        self.assertIn(wn_in_value.get(instance=i), prim_ids)
                        self.assertIn(wn_out_value.get(instance=i), prim_ids)
                        self.assertEqual(out_var.get(graph_context, instance_path=prim_paths[i]), expected_value)
                        if (
                            evaluator_name == "push"
                        ):  # Needed because PushGraphs have some special vectorization going on?
                            if si == 0:
                                self.assertEqual(rn_cnt.get(), 2)
                                self.assertEqual(wn_cnt.get(), 2)
                                self.assertEqual(rn_cnt.get(instance=i), 0)
                                self.assertEqual(wn_cnt.get(instance=i), 0)
                            else:
                                self.assertEqual(rn_cnt.get(), 0)
                                self.assertEqual(wn_cnt.get(), 0)
                                if i == 0:
                                    self.assertEqual(rn_cnt.get(instance=i), 2)
                                    self.assertEqual(wn_cnt.get(instance=i), 2)
                                else:
                                    self.assertEqual(rn_cnt.get(instance=i), 0)
                                    self.assertEqual(wn_cnt.get(instance=i), 0)
                        else:
                            if si == 0:
                                self.assertEqual(rn_cnt.get(), 6)
                                self.assertEqual(wn_cnt.get(), 6)
                                self.assertEqual(rn_cnt.get(instance=i), 0)
                                self.assertEqual(wn_cnt.get(instance=i), 0)
                            else:
                                self.assertEqual(rn_cnt.get(), 0)
                                self.assertEqual(wn_cnt.get(), 0)
                                if i == 0:
                                    self.assertEqual(rn_cnt.get(instance=i), 6)
                                    self.assertEqual(wn_cnt.get(instance=i), 6)
                                else:
                                    self.assertEqual(rn_cnt.get(instance=i), 0)
                                    self.assertEqual(wn_cnt.get(instance=i), 0)

                    # Delete all prims in the stage.
                    for i in range(0, num_prims):
                        stage.RemovePrim(prim_paths[i])
                    graph.remove_variable(in_var)
                    graph.remove_variable(out_var)
                    stage.RemovePrim(graph_path)
                    await omni.kit.app.get_app().next_update_async()

    async def test_constant_node_removal(self):
        """
        Check that we do not include constant OG nodes in the execution
        graph backend for push and action graphs, and that doing so does
        not impact graph computation to yield incorrect results.
        """

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

        # Run the test over each evaluator type.
        evaluator_names = [
            "push",
            "dirty_push",
            "execution",
        ]

        for evaluator_name in evaluator_names:
            # Create our test graph.
            graph_path = "/World/TestGraph"
            test_prim_path = "/World/TestPrim"
            (graph, (constant_double0_node, constant_double1_node, add_node, _, _, _), _, _) = og.Controller.edit(
                {"graph_path": graph_path, "evaluator_name": evaluator_name},
                {
                    og.Controller.Keys.CREATE_NODES: [
                        ("ConstantDouble0", "omni.graph.nodes.ConstantDouble"),
                        ("ConstantDouble1", "omni.graph.nodes.ConstantDouble"),
                        ("Add", "omni.graph.nodes.Add"),
                        ("OnTick", "omni.graph.action.OnTick"),
                        ("WritePrimAttribute", "omni.graph.nodes.WritePrimAttribute"),
                        ("WriteVariable", "omni.graph.core.WriteVariable"),
                    ],
                    og.Controller.Keys.CREATE_PRIMS: (
                        test_prim_path,
                        {
                            "value": ("double", 0),
                        },
                    ),
                    og.Controller.Keys.CREATE_VARIABLES: [
                        ("out_double_var", og.Type(og.BaseDataType.DOUBLE)),
                    ],
                    og.Controller.Keys.CONNECT: [
                        ("ConstantDouble0.inputs:value", "Add.inputs:a"),
                        ("ConstantDouble1.inputs:value", "Add.inputs:b"),
                        ("Add.outputs:sum", "WritePrimAttribute.inputs:value"),
                        ("Add.outputs:sum", "WriteVariable.inputs:value"),
                        ("OnTick.outputs:tick", "WritePrimAttribute.inputs:execIn"),
                        ("WritePrimAttribute.outputs:execOut", "WriteVariable.inputs:execIn"),
                    ],
                    og.Controller.Keys.SET_VALUES: [
                        ("ConstantDouble0.inputs:value", 5.0),
                        ("ConstantDouble1.inputs:value", 1.0),
                        ("OnTick.inputs:onlyPlayback", False),
                        ("WritePrimAttribute.inputs:name", "value"),
                        ("WritePrimAttribute.inputs:primPath", "/World/TestPrim"),
                        ("WriteVariable.inputs:variableName", "out_double_var"),
                    ],
                },
            )

            out_double_var = graph.find_variable("out_double_var")
            inputs_value_attr_0 = constant_double0_node.get_attribute("inputs:value")
            inputs_value_attr_1 = constant_double1_node.get_attribute("inputs:value")
            outputs_sum_attr = add_node.get_attribute("outputs:sum")

            # Check that the summation is correctly computed, and that the
            # ConstantDouble nodes aren't computed.
            await omni.kit.app.get_app().next_update_async()
            self.assertAlmostEqual(outputs_sum_attr.get(), 6.0, places=2)
            if evaluator_name in ("push", "execution"):
                self.assertEqual(constant_double0_node.get_compute_count(), 0)
                self.assertEqual(constant_double1_node.get_compute_count(), 0)
            else:
                self.assertEqual(constant_double0_node.get_compute_count(), 1)
                self.assertEqual(constant_double1_node.get_compute_count(), 1)

            # Check that this still holds if we alter attributes on the nodes.
            inputs_value_attr_0.set(8)
            inputs_value_attr_1.set(2)
            await omni.kit.app.get_app().next_update_async()
            self.assertAlmostEqual(outputs_sum_attr.get(), 10.0, places=2)
            if evaluator_name in ("push", "execution"):
                self.assertEqual(constant_double0_node.get_compute_count(), 0)
                self.assertEqual(constant_double1_node.get_compute_count(), 0)
            else:
                self.assertEqual(constant_double0_node.get_compute_count(), 2)
                self.assertEqual(constant_double1_node.get_compute_count(), 2)

            # Next, instance the graph 3 times and perform similar checks as before.
            graph_context = graph.get_default_graph_context()
            num_graph_instances = 3
            graph_instance_ids = [None] * num_graph_instances
            graph_instance_paths = [None] * num_graph_instances
            graph_instances = [None] * num_graph_instances
            for i in range(0, num_graph_instances):
                graph_instance_ids[i] = i
                graph_instance_paths[i] = f"/World/GraphInstance_{i}"
                graph_instances[i] = stage.DefinePrim(graph_instance_paths[i])
                OmniGraphSchemaTools.applyOmniGraphAPI(stage, graph_instance_paths[i], graph_path)
                graph_instances[i].CreateAttribute("graph:variable:out_double_var", Sdf.ValueTypeNames.Double).Set(0)

            inputs_value_attr_0.set(1)
            inputs_value_attr_1.set(3)
            await omni.kit.app.get_app().next_update_async()
            for i in range(0, num_graph_instances):
                # TODO: Instanced "execution" type graphs aren't getting updated here
                # for some reason, needs a bit more investigation...
                if evaluator_name == "execution":
                    self.assertEqual(out_double_var.get(graph_context, instance_path=graph_instance_paths[i]), 0.0)
                else:
                    self.assertEqual(out_double_var.get(graph_context, instance_path=graph_instance_paths[i]), 4.0)
            if evaluator_name in ("push", "execution"):
                self.assertEqual(constant_double0_node.get_compute_count(), 0)
                self.assertEqual(constant_double1_node.get_compute_count(), 0)
            else:
                self.assertEqual(constant_double0_node.get_compute_count(), 6)
                self.assertEqual(constant_double1_node.get_compute_count(), 6)

            # Delete all prims in the stage.
            for i in range(0, num_graph_instances):
                stage.RemovePrim(graph_instance_paths[i])
            stage.RemovePrim(test_prim_path)
            stage.RemovePrim(graph_path)
            await omni.kit.app.get_app().next_update_async()
