# noqa PLC0302
"""Tests that exercise the functionality of the extended (union and any) attribute types"""
import os
import tempfile
from typing import List, Optional

import carb
import numpy as np
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.app
import omni.kit.stage_templates
import omni.kit.test
import omni.usd
from pxr import Gf, Sdf, UsdGeom, Vt


# ======================================================================
class TestExtendedAttributes(ogts.OmniGraphTestCase):
    """Run a simple unit test that exercises graph functionality"""

    TEST_GRAPH_PATH = "/World/TestGraph"

    # ----------------------------------------------------------------------
    async def test_output_type_resolution(self):
        """Test the automatic type resolution for extended types based on output connection types"""
        controller = og.Controller()
        keys = og.Controller.Keys

        # Node is by itself, union attributes have no type (use Attribute.getTypeName to check)
        (graph, (extended_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH, {keys.CREATE_NODES: ("ExtendedNode", "omni.graph.test.TypeResolution")}
        )
        value_output = "outputs:value"
        array_output = "outputs:arrayValue"
        tuple_output = "outputs:tupleValue"
        tuple_array_output = "outputs:tupleArrayValue"
        mixed_output = "outputs:mixedValue"
        any_output = "outputs:anyValue"
        resolved_type = controller.attribute("outputs:resolvedType", extended_node)
        value = controller.attribute(value_output, extended_node)
        array_value = controller.attribute(array_output, extended_node)
        tuple_value = controller.attribute(tuple_output, extended_node)
        tuple_array_value = controller.attribute(tuple_array_output, extended_node)
        mixed_value = controller.attribute(mixed_output, extended_node)
        any_value = controller.attribute(any_output, extended_node)
        await controller.evaluate(graph)

        def result():
            return og.Controller.get(resolved_type)

        def expected(
            value_type: Optional[str] = None,
            array_type: Optional[str] = None,
            tuple_type: Optional[str] = None,
            tuple_array_type: Optional[str] = None,
            mixed_type: Optional[str] = None,
            any_type: Optional[str] = None,
        ) -> List[str]:
            """Return an expected output array for the given resolved types (None means unresolved)"""
            return [
                f"{value_output},{value_type if value_type else 'unknown'}",
                f"{array_output},{array_type if array_type else 'unknown'}",
                f"{tuple_output},{tuple_type if tuple_type else 'unknown'}",
                f"{tuple_array_output},{tuple_array_type if tuple_array_type else 'unknown'}",
                f"{mixed_output},{mixed_type if mixed_type else 'unknown'}",
                f"{any_output},{any_type if any_type else 'unknown'}",
            ]

        default_expected = expected()
        default_actual = result()
        self.assertEqual(len(default_expected), len(default_actual))
        self.assertCountEqual(default_expected, default_actual)

        # Node has float/token output connected to a float input - type should be float
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("Simple", "omni.graph.tutorials.SimpleData"),
                keys.CONNECT: (value, "Simple.inputs:a_float"),
            },
        )
        await controller.evaluate(graph)
        self.assertCountEqual(expected("float"), result())

        # Node disconnects, then connects the same output to an int input - type should be int
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.DISCONNECT: (value, "Simple.inputs:a_float"),
            },
        )
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CONNECT: (value, "Simple.inputs:a_int"),
            },
        )
        await controller.evaluate(graph)
        self.assertCountEqual(expected("int"), result())

        # Node totally disconnects - type should be reverted to unknown
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.DISCONNECT: (value, "Simple.inputs:a_int"),
            },
        )
        await controller.evaluate(graph)
        self.assertCountEqual(expected(), result())

        # Adding array type connection
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("Array", "omni.graph.tutorials.ArrayData"),
                keys.CONNECT: (array_value, "Array.inputs:original"),
            },
        )
        await controller.evaluate(graph)
        self.assertCountEqual(expected("unknown", "float[]"), result())

        # Adding tuple, tupleArray, and mixed type connections
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Tuple", "omni.tutorials.TupleData"),
                    ("TupleArray", "omni.graph.tutorials.TupleArrays"),
                ],
                keys.CONNECT: [
                    (tuple_value, "Tuple.inputs:a_float3"),
                    (tuple_array_value, "TupleArray.inputs:a"),
                    (mixed_value, "Array.inputs:original"),
                ],
            },
        )
        await controller.evaluate(graph)
        self.assertCountEqual(expected("unknown", "float[]", "float[3]", "float[3][]", "float[]"), result())

        # Adding anyValue connection, check type
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CONNECT: (any_value, "Simple.inputs:a_float"),
            },
        )
        await controller.evaluate(graph)
        self.assertCountEqual(expected("unknown", "float[]", "float[3]", "float[3][]", "float[]", "float"), result())

        # Now test the type-resolution cascade for a linear chain of any-type connections
        (_, (extended_node_a, extended_node_b, extended_node_c), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("ExtendedNodeA", "omni.graph.test.TypeResolution"),
                    ("ExtendedNodeB", "omni.graph.test.TypeResolution"),
                    ("ExtendedNodeC", "omni.graph.test.TypeResolution"),
                ],
                keys.CONNECT: [
                    ("ExtendedNodeA.outputs:anyValue", "ExtendedNodeB.inputs:anyValueIn"),
                    ("ExtendedNodeB.outputs:anyValue", "ExtendedNodeC.inputs:anyValueIn"),
                ],
            },
        )
        await controller.evaluate(graph)

        node_a_controller = og.Controller(og.Controller.attribute("outputs:resolvedType", extended_node_a))
        node_b_controller = og.Controller(og.Controller.attribute("outputs:resolvedType", extended_node_b))
        node_c_controller = og.Controller(og.Controller.attribute("outputs:resolvedType", extended_node_c))

        resolved_types = node_c_controller.get()
        self.assertCountEqual(expected(), resolved_types)

        # connect a concrete type to the start of the chain and verify it cascades to the last output
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CONNECT: ("Simple.outputs:a_double", "ExtendedNodeA.inputs:anyValueIn"),
            },
        )
        await controller.evaluate(graph)

        resolved_types = node_a_controller.get()
        self.assertCountEqual(expected(any_type="double"), resolved_types)
        resolved_types = node_b_controller.get()
        self.assertCountEqual(expected(any_type="double"), resolved_types)
        resolved_types = node_c_controller.get()
        self.assertCountEqual(expected(any_type="double"), resolved_types)

        # remove the concrete connection and verify it cascades unresolved to the last output
        # FIXME: Auto-unresolve is not supported yet
        # controller.edit(self.TEST_GRAPH_PATH, {
        #     keys.DISCONNECT: ("Simple.outputs:a_double", "ExtendedNodeA.inputs:anyValueIn"),
        # })
        # await controller.evaluate(graph)
        #
        # resolved_types = node_a_controller.get()
        # self.assertCountEqual(expected(), resolved_types)
        # resolved_types = node_b_controller.get()
        # self.assertCountEqual(expected(), resolved_types)
        # resolved_types = node_c_controller.get()
        # self.assertCountEqual(expected(), resolved_types)

        # add a concrete type to the output connection at the end of the chain and verify it cascades upwards
        # to the input of the start of the chain
        (_, (extended_node_a, extended_node_b, extended_node_c), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.DELETE_NODES: [
                    extended_node_a,
                    extended_node_b,
                    extended_node_c,
                ],
                keys.CREATE_NODES: [
                    ("ExtendedNodeA", "omni.graph.test.TypeResolution"),
                    ("ExtendedNodeB", "omni.graph.test.TypeResolution"),
                    ("ExtendedNodeC", "omni.graph.test.TypeResolution"),
                ],
                keys.CONNECT: [
                    ("ExtendedNodeA.outputs:anyValue", "ExtendedNodeB.inputs:anyValueIn"),
                    ("ExtendedNodeB.outputs:anyValue", "ExtendedNodeC.inputs:anyValueIn"),
                    ("ExtendedNodeC.outputs:anyValue", "Simple.inputs:a_double"),
                ],
            },
        )
        await controller.evaluate(graph)

        node_a_controller = og.Controller(og.Controller.attribute("outputs:resolvedType", extended_node_a))
        node_b_controller = og.Controller(og.Controller.attribute("outputs:resolvedType", extended_node_b))
        node_c_controller = og.Controller(og.Controller.attribute("outputs:resolvedType", extended_node_c))

        resolved_types = node_c_controller.get()
        self.assertCountEqual(expected(any_type="double"), resolved_types)
        resolved_types = node_b_controller.get()
        self.assertCountEqual(expected(any_type="double"), resolved_types)
        resolved_types = node_a_controller.get()
        self.assertCountEqual(expected(any_type="double"), resolved_types)

    # ----------------------------------------------------------------------
    async def test_type_propagation(self):
        """Test the propagation of type resolution through the network"""
        controller = og.Controller()
        keys = og.Controller.Keys

        # Test with a more complicated diamond-shaped topology
        #                             +-----+
        #                         +-->|  B  +-+   +-----+
        #     +-----+    +------+-+   +-----+ |-->|     |   +-----+
        #     |  SA +--->|   A  |                 | Add |-->|  SB +
        #     +-----+    +------+-+   +-----+ |-->|     |   +-----+
        #                         +-->|  C  +-+   +-----+
        #                             +-----+
        #

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("SimpleA", "omni.graph.tutorials.SimpleData"),
                    ("ExtendedNodeA", "omni.graph.test.TypeResolution"),
                    ("ExtendedNodeB", "omni.graph.test.TypeResolution"),
                    ("ExtendedNodeC", "omni.graph.test.TypeResolution"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("SimpleB", "omni.graph.tutorials.SimpleData"),
                ],
                keys.CONNECT: [
                    ("ExtendedNodeA.outputs:anyValue", "ExtendedNodeB.inputs:anyValueIn"),
                    ("ExtendedNodeA.outputs:anyValue", "ExtendedNodeC.inputs:anyValueIn"),
                    ("ExtendedNodeB.outputs:anyValue", "Add.inputs:a"),
                    ("ExtendedNodeC.outputs:anyValue", "Add.inputs:b"),
                    ("Add.outputs:sum", "SimpleB.inputs:a_double"),
                    # Concrete double wired to front of chain to propagate forward
                    ("SimpleA.outputs:a_double", "ExtendedNodeA.inputs:anyValueIn"),
                ],
                keys.SET_VALUES: [
                    ("SimpleA.inputs:a_double", 10.0),
                ],
            },
        )
        await controller.evaluate(graph)

        (simple_a_node, extended_node_a, _, _, add_node, simple_b_node) = nodes
        v = og.Controller.get(controller.attribute("outputs:a_double", simple_a_node))
        self.assertEqual(v, 11.0)
        v = og.Controller.get(controller.attribute("outputs:anyValue", extended_node_a))
        self.assertEqual(v, 11.0)
        v = og.Controller.get(controller.attribute("outputs:sum", add_node))
        self.assertEqual(v, 22.0)

        # connect the last any output to a concrete input, verify value
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CONNECT: ("Add.outputs:sum", "SimpleB.inputs:a_double"),
            },
        )
        await controller.evaluate(graph)

        v = og.Controller.get(controller.attribute("outputs:a_double", simple_b_node))
        self.assertEqual(v, 23.0)

        # disconnect and re-connect the connection, verify the value is the same (tests output re-connect)
        controller.edit(self.TEST_GRAPH_PATH, {keys.DISCONNECT: ("Add.outputs:sum", "SimpleB.inputs:a_double")})
        await controller.evaluate(graph)

        controller.edit(self.TEST_GRAPH_PATH, {keys.CONNECT: ("Add.outputs:sum", "SimpleB.inputs:a_double")})
        await controller.evaluate(graph)
        v = og.Controller.get(controller.attribute("outputs:a_double", simple_b_node))
        self.assertEqual(v, 23.0)

        # disconnect and re-connect the extended inputs to sum and verify (tests input re-connect)
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.DISCONNECT: [
                    ("ExtendedNodeB.outputs:anyValue", "Add.inputs:a"),
                    ("ExtendedNodeC.outputs:anyValue", "Add.inputs:b"),
                ]
            },
        )
        await controller.evaluate(graph)
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CONNECT: [
                    ("ExtendedNodeB.outputs:anyValue", "Add.inputs:a"),
                    ("ExtendedNodeC.outputs:anyValue", "Add.inputs:b"),
                ]
            },
        )
        await controller.evaluate(graph)
        v = og.Controller.get(controller.attribute("outputs:a_double", simple_b_node))
        self.assertEqual(v, 23.0)

    # ----------------------------------------------------------------------
    async def test_sandwiched_type_resolution(self):
        """modulo (which only operate on integers) is sandwiched between 2 nodes that feed/expect double
        this setup will test that this conversion are correctly chosen"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Modulo", "omni.graph.nodes.Modulo"),
                    ("ConstD", "omni.graph.nodes.ConstantDouble"),
                    ("ConstF", "omni.graph.nodes.ConstantFloat"),
                    ("Scale", "omni.graph.nodes.ScaleToSize"),
                ],
                keys.CONNECT: [
                    ("ConstD.inputs:value", "Modulo.inputs:a"),
                    ("ConstD.inputs:value", "Modulo.inputs:b"),
                    ("Modulo.outputs:result", "Scale.inputs:speed"),
                ],
            },
        )
        await controller.evaluate(graph)

        # Helper to verify the base data type of an attribute
        def assert_attrib_is(attrib, node, base_type):
            attrib = controller.attribute(attrib, node, graph)
            self.assertEqual(attrib.get_resolved_type().base_type, base_type)

        assert_attrib_is("outputs:result", "Modulo", og.BaseDataType.INT64)

        # Disconnect the Multiply inputs to cause a wave of unresolutions
        controller.edit(
            graph,
            {keys.DISCONNECT: [("ConstD.inputs:value", "Modulo.inputs:a"), ("ConstD.inputs:value", "Modulo.inputs:b")]},
        )
        await controller.evaluate(graph)
        assert_attrib_is("outputs:result", "Modulo", og.BaseDataType.UNKNOWN)

        # now connect a float on input:a, that should make it the prefered type to choose conversion from
        #  which will be INT this time
        controller.edit(
            graph,
            {keys.CONNECT: [("ConstF.inputs:value", "Modulo.inputs:a"), ("ConstD.inputs:value", "Modulo.inputs:b")]},
        )
        await controller.evaluate(graph)
        assert_attrib_is("outputs:result", "Modulo", og.BaseDataType.INT)

    # ----------------------------------------------------------------------
    async def test_reconnect_different_type(self):
        """Tests re-resolving and the RESOLVE_ATTRIBUTE event"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (_, node, _, _), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("ForEach", "omni.graph.action.ForEach"),
                    ("Get", "omni.graph.nodes.ReadPrimAttribute"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: [
                    ("/World/Source", {"myfloatarray": ("float[]", [1000.5]), "myintarray": ("int64[]", [2**40])}),
                    ("/World/Sink", {"myfloat": ("float", 0), "myint": ("int64", 0)}),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "ForEach.inputs:execIn"),
                    ("Get.outputs:value", "ForEach.inputs:arrayIn"),
                    ("ForEach.outputs:element", "Set.inputs:value"),
                ],
                keys.SET_VALUES: [
                    ("Set.inputs:name", "myfloat"),
                    ("Set.inputs:primPath", "/World/Sink"),
                    ("Set.inputs:usePath", True),
                    ("Get.inputs:name", "myfloatarray"),
                    ("Get.inputs:primPath", "/World/Source"),
                    ("Get.inputs:usePath", True),
                ],
            },
        )
        await controller.evaluate(graph)
        v = og.Controller.get(controller.attribute("outputs:element", node))
        self.assertEqual(v, 1000.5)

        # Trigger unresolve and re-resolve of the 2 ForEach attributes
        event_counter = [0]

        def on_node_event(event):
            self.assertEqual(event.type, int(og.NodeEvent.ATTRIBUTE_TYPE_RESOLVE))
            event_counter[0] += 1

        # Hold on to the sub!
        sub = node.get_event_stream().create_subscription_to_pop(on_node_event)

        with ogts.ExpectedError():
            controller.edit(
                self.TEST_GRAPH_PATH,
                {keys.SET_VALUES: [("Set.inputs:name", "myint"), ("Get.inputs:name", "myintarray")]},
            )
        await controller.evaluate(graph)
        await controller.evaluate(graph)
        v = og.Controller.get(controller.attribute("outputs:element", node))
        self.assertEqual(v, 2**40)

        # Check we got the right number of callbacks (2x callbacks for 2 attributes)
        self.assertEqual(event_counter[0], 4)

        del sub

    # ----------------------------------------------------------------------
    # Tests of ReadPrimAttribute follow

    async def test_readprimattribute_scaler(self):
        """Exercise ReadPrimAttribute for scaler value"""

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

        (graph, (get_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("Get", "omni.graph.nodes.ReadPrimAttribute"),
                keys.CREATE_PRIMS: ("/World/Float", {"myfloat": ("float", 42.0)}),
                keys.SET_VALUES: [
                    ("Get.inputs:name", "myfloat"),
                    ("Get.inputs:primPath", "/World/Float"),
                    ("Get.inputs:usePath", True),
                ],
            },
        )
        await controller.evaluate(graph)
        await controller.evaluate(graph)
        get_node_controller = og.Controller(og.Controller.attribute("outputs:value", get_node))
        self.assertEqual(get_node_controller.get(), 42.0)
        # Change USD attrib, check we get the new value on our output
        attr = stage.GetPropertyAtPath("/World/Float.myfloat")
        attr.Set(0.0)
        await controller.evaluate(graph)
        self.assertEqual(get_node_controller.get(), 0.0)

        # Same test but with bundle input

        stage.GetPrimAtPath(f"{self.TEST_GRAPH_PATH}/Get").GetRelationship("inputs:prim").SetTargets(
            [Sdf.Path("/World/Float")]
        )
        attr.Set(42.0)
        controller.edit(self.TEST_GRAPH_PATH, {keys.SET_VALUES: ("Get.inputs:usePath", False)})
        await controller.evaluate(graph)
        self.assertEqual(get_node_controller.get(), 42.0)
        # Change USD attrib, check we get the new value on our output
        attr = stage.GetPropertyAtPath("/World/Float.myfloat")
        attr.Set(0.0)
        await controller.evaluate(graph)
        self.assertEqual(get_node_controller.get(), 0.0)

    # ----------------------------------------------------------------------
    async def test_readprimattribute_noexist(self):
        """Test ReadPrimAttribute can handle attribute that exists in USD but not in FC"""

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

        (graph, (get_node,), (prim,), _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("Get", "omni.graph.nodes.ReadPrimAttribute"),
                keys.CREATE_PRIMS: ("/World/Cube", "Cube"),
                keys.SET_VALUES: [
                    ("Get.inputs:name", "visibility"),
                    ("Get.inputs:primPath", "/World/Cube"),
                    ("Get.inputs:usePath", True),
                ],
            },
        )
        await controller.evaluate(graph)
        await controller.evaluate(graph)

        self.assertEqual(og.Controller.get(controller.attribute("outputs:value", get_node)), "inherited")

        # Change USD attrib, check we get the new value on our output
        UsdGeom.Imageable(prim).MakeInvisible()

        await controller.evaluate(graph)
        # ReadPrimAttribute will put the visibility value to Fabric, but it doesn't get
        # updated by the FSD USD handler.
        # Instead it updates the visibility of the prim, so this is a workaround to test
        # that the visibility of the prim been updated.
        if carb.settings.get_settings().get("/app/useFabricSceneDelegate"):
            import usdrt

            stage = usdrt.Usd.Stage.Attach(omni.usd.get_context().get_stage_id())
            prim = stage.GetPrimAtPath("/World/Cube")
            attr = prim.GetAttribute("_worldVisibility")
            self.assertFalse(attr.Get())
        else:
            self.assertEqual(og.Controller.get(controller.attribute("outputs:value", get_node)), "invisible")

    # ----------------------------------------------------------------------
    async def test_readprimattribute_tuple(self):
        """Exercise ReadPrimAttribute for a tuple value"""

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

        (graph, (get_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("Get", "omni.graph.nodes.ReadPrimAttribute"),
                keys.CREATE_PRIMS: ("/World/Double3", {"mydouble3": ("double[3]", Gf.Vec3d(42.0, 40.0, 1.0))}),
                keys.SET_VALUES: [
                    ("Get.inputs:name", "mydouble3"),
                    ("Get.inputs:primPath", "/World/Double3"),
                    ("Get.inputs:usePath", True),
                ],
            },
        )
        await controller.evaluate(graph)
        og.Controller.get(controller.attribute("outputs:value", get_node))
        self.assertEqual(
            Gf.Vec3d(*og.Controller.get(controller.attribute("outputs:value", get_node))), Gf.Vec3d(42.0, 40.0, 1.0)
        )
        # Change USD attrib, check we get the new value on our output
        attr = stage.GetPropertyAtPath("/World/Double3.mydouble3")
        attr.Set(Gf.Vec3d(0, 0, 0))
        await controller.evaluate(graph)
        self.assertEqual(
            Gf.Vec3d(*og.Controller.get(controller.attribute("outputs:value", get_node))), Gf.Vec3d(0, 0, 0)
        )

    # ----------------------------------------------------------------------
    async def test_readprimattribute_array(self):
        """Exercise ReadPrimAttribute for an array value"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        small_array = [0xFF for _ in range(10)]

        (graph, (get_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: ("Get", "omni.graph.nodes.ReadPrimAttribute"),
                keys.CREATE_PRIMS: ("/World/IntArray", {"myintarray": ("int[]", small_array)}),
                keys.SET_VALUES: [
                    ("Get.inputs:name", "myintarray"),
                    ("Get.inputs:usePath", True),
                    ("Get.inputs:primPath", "/World/IntArray"),
                ],
            },
        )
        await controller.evaluate(graph)
        self.assertSequenceEqual(
            og.Controller.get(controller.attribute("outputs:value", get_node)).tolist(), small_array
        )
        # Change USD attrib, check we get the new value on our output
        big_array = Vt.IntArray(1000, [0 for _ in range(1000)])
        attr = stage.GetPropertyAtPath("/World/IntArray.myintarray")
        attr.Set(big_array)
        await controller.evaluate(graph)
        self.assertSequenceEqual(og.Controller.get(controller.attribute("outputs:value", get_node)).tolist(), big_array)

    # ----------------------------------------------------------------------
    async def test_dynamic_extended_attributes_any(self):
        """Test functionality of extended any attributes when they are created dynamically"""

        # Set up a graph that tests dynamically created extended attrs.  Here we have a node with
        # extended attribute attr1.  We create another extended attribute dynamically (dynamic_attr1)
        # we then make sure these two behave the same. dyanmic_attr1 is an any attribute in this case
        #
        #    SimpleIn ----=> Extended1
        #                       -inputs:anyValueIn
        #                       -inputs:dynamic_attr1
        #
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (simple_node, extended_node_1), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("SimpleIn", "omni.graph.tutorials.SimpleData"),
                    ("Extended1", "omni.graph.test.TypeResolution"),
                ],
                keys.CONNECT: ("SimpleIn.outputs:a_float", "Extended1.inputs:anyValueIn"),
                keys.SET_VALUES: ("SimpleIn.inputs:a_float", 5.0),
            },
        )
        await controller.evaluate(graph)

        # Check that the inputs into the extended type nodes are correct
        self.assertEqual(6.0, og.Controller.get(controller.attribute("inputs:anyValueIn", extended_node_1)))
        extended_node_1.create_attribute(
            "inputs:dynamic_attr1",
            og.Type(og.BaseDataType.TOKEN),
            og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
            None,
            og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_ANY,
            "",
        )

        controller.edit(
            self.TEST_GRAPH_PATH,
            {keys.CONNECT: (("outputs:a_float", simple_node), ("inputs:dynamic_attr1", extended_node_1))},
        )
        await controller.evaluate(graph)
        self.assertEqual(6.0, og.Controller.get(controller.attribute("inputs:dynamic_attr1", extended_node_1)))

    # ----------------------------------------------------------------------
    async def test_dynamic_extended_attributes_union(self):
        """Test functionality of extended union attributes when they are created dynamically"""

        # Set up a graph that tests dynamically created extended attrs.  Here we have a node with
        # extended attribute attr1.  We create another extended attribute dynamically (dynamic_attr1)
        # we then make sure these two behave the same.  dyanmic_attr1 is a union attribute in this case
        #
        #    SimpleIn ----=> Extended1
        #                       -inputs:anyValueIn
        #                       -inputs:dynamic_attr1
        #
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (simple_node, extended_node_1), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("SimpleIn", "omni.graph.tutorials.SimpleData"),
                    ("Extended1", "omni.graph.tutorials.ExtendedTypes"),
                ],
                keys.CONNECT: ("SimpleIn.outputs:a_float", "Extended1.inputs:floatOrToken"),
                keys.SET_VALUES: ("SimpleIn.inputs:a_float", 5.0),
            },
        )
        await controller.evaluate(graph)

        # Check that the inputs into the extended type nodes are correct
        self.assertEqual(6.0, og.Controller.get(controller.attribute("inputs:floatOrToken", extended_node_1)))
        extended_node_1.create_attribute(
            "inputs:dynamic_attr1",
            og.Type(og.BaseDataType.TOKEN),
            og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
            None,
            og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_UNION,
            "float,token",
        )
        controller.edit(
            self.TEST_GRAPH_PATH,
            {keys.CONNECT: (("outputs:a_float", simple_node), ("inputs:dynamic_attr1", extended_node_1))},
        )
        await controller.evaluate(graph)
        self.assertEqual(6.0, og.Controller.get(controller.attribute("inputs:dynamic_attr1", extended_node_1)))

    # ----------------------------------------------------------------------
    async def test_unresolve_propagation(self):
        """Tests unresolve propagates through the network"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (_, _, _, _, to_string_node, _, _), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("ConstDouble", "omni.graph.nodes.ConstantDouble"),
                    ("Const1", "omni.graph.nodes.ConstantVec3d"),
                    ("Multiply", "omni.graph.nodes.Multiply"),
                    ("Magnitude", "omni.graph.nodes.Magnitude"),
                    ("ToString", "omni.graph.nodes.ToString"),
                    ("Const2", "omni.graph.nodes.ConstantDouble"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: [("/World/Sink", {"mydouble": ("double", 0)})],
                keys.SET_VALUES: [
                    ("ConstDouble.inputs:value", 2.0),
                    ("Const1.inputs:value", (0.5, 0, 0)),
                    ("Set.inputs:name", "mydouble"),
                    ("Set.inputs:primPath", "/World/Sink"),
                    ("Set.inputs:usePath", True),
                ],
                keys.CONNECT: [
                    ("ConstDouble.inputs:value", "Multiply.inputs:a"),
                    ("Const1.inputs:value", "Multiply.inputs:b"),
                    ("Multiply.outputs:product", "Magnitude.inputs:input"),
                    ("Magnitude.outputs:magnitude", "ToString.inputs:value"),
                ],
            },
        )
        await controller.evaluate(graph)
        v = og.Controller.get(controller.attribute("inputs:value", to_string_node))
        self.assertEqual(v, 1.0)
        # Disconnect one of the Multiply inputs to cause a wave of unresolutions
        controller.edit(
            graph,
            {
                keys.DISCONNECT: [
                    ("ConstDouble.inputs:value", "Multiply.inputs:a"),
                    ("Const1.inputs:value", "Multiply.inputs:b"),
                ]
            },
        )

        # Helper to verify the base data type of an attribute
        def assert_attrib_is(attrib, node, base_type):
            attrib = controller.attribute(attrib, node, graph)
            self.assertEqual(attrib.get_resolved_type().base_type, base_type)

        # Verify the irresolute quality of the downstream attributes
        assert_attrib_is("outputs:product", "Multiply", og.BaseDataType.UNKNOWN)
        assert_attrib_is("inputs:input", "Magnitude", og.BaseDataType.UNKNOWN)
        assert_attrib_is("inputs:value", "ToString", og.BaseDataType.UNKNOWN)

        # Re-connect to observe resolution
        controller.edit(
            graph,
            {
                keys.CONNECT: [
                    ("ConstDouble.inputs:value", "Multiply.inputs:a"),
                    ("Const1.inputs:value", "Multiply.inputs:b"),
                ]
            },
        )
        # Verify the resolute quality of the downstream attributes
        assert_attrib_is("outputs:product", "Multiply", og.BaseDataType.DOUBLE)
        assert_attrib_is("inputs:input", "Magnitude", og.BaseDataType.DOUBLE)
        assert_attrib_is("inputs:value", "ToString", og.BaseDataType.DOUBLE)

        # Disconnect as before
        controller.edit(
            graph,
            {
                keys.DISCONNECT: [
                    ("ConstDouble.inputs:value", "Multiply.inputs:a"),
                    ("Const1.inputs:value", "Multiply.inputs:b"),
                ]
            },
        )
        # Wait on the next framework tick, we don't need a full evaluation
        await omni.kit.app.get_app().next_update_async()

        # Verify that this time the unresolution was blocked at Magnitude
        assert_attrib_is("outputs:product", "Multiply", og.BaseDataType.UNKNOWN)
        assert_attrib_is("inputs:a", "Multiply", og.BaseDataType.UNKNOWN)
        assert_attrib_is("inputs:input", "Magnitude", og.BaseDataType.UNKNOWN)
        assert_attrib_is("inputs:value", "ToString", og.BaseDataType.UNKNOWN)

        # Now check that WritePrimAttribute will cause upstream propagation
        await controller.evaluate(graph)
        controller.edit(graph, {keys.CONNECT: [("Magnitude.outputs:magnitude", "Set.inputs:value")]})
        await controller.evaluate(graph)
        # Verify that our network has flipped over to double
        assert_attrib_is("outputs:magnitude", "Magnitude", og.BaseDataType.DOUBLE)
        # Note that we no longer propagate resolution upstream from outputs->inputs, so unresolution
        # remains upstream of Magnitude
        assert_attrib_is("inputs:input", "Magnitude", og.BaseDataType.UNKNOWN)
        assert_attrib_is("outputs:product", "Multiply", og.BaseDataType.UNKNOWN)
        assert_attrib_is("inputs:a", "Multiply", og.BaseDataType.UNKNOWN)
        assert_attrib_is("inputs:value", "ToString", og.BaseDataType.DOUBLE)

    # ----------------------------------------------------------------------
    async def test_unresolve_contradiction(self):
        """Tests unresolve contradiction is handled"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("ConstVec3f", "omni.graph.nodes.ConstantVec3f"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CONNECT: [
                    ("Add.outputs:sum", "Set.inputs:value"),
                ],
            },
        )
        await omni.kit.app.get_app().next_update_async()

        event_counter = [0]

        def on_node_event(event):
            event_counter[0] += 1
            self.assertEqual(event.type, int(og.NodeEvent.ATTRIBUTE_TYPE_RESOLVE))

        node = controller.node("Add", graph)
        sub = node.get_event_stream().create_subscription_to_pop(on_node_event)

        controller.edit(
            graph,
            {
                keys.CONNECT: [
                    ("ConstVec3f.inputs:value", "Add.inputs:a"),
                    ("ConstVec3f.inputs:value", "Add.inputs:b"),
                ]
            },
        )
        for _ in range(10):
            await omni.kit.app.get_app().next_update_async()
        self.assertLess(event_counter[0], 10)

        del sub

    # ----------------------------------------------------------------------
    async def test_unresolve_resolve_pingpong(self):
        """Tests resolve doesn't ping-pong is handled"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (_, get_node, _), (p_in, p_out), _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Get", "omni.graph.nodes.ReadPrimAttribute"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: [("PIn", {"f": ("float", 1.0)}), ("POut", {"f": ("float", 1.0)})],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "Set.inputs:execIn"),
                    ("Get.outputs:value", "Set.inputs:value"),
                ],
            },
        )
        controller.edit(
            graph,
            {
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("Get.inputs:usePath", True),
                    ("Get.inputs:usePath", True),
                    ("Get.inputs:name", "f"),
                    ("Set.inputs:name", "f"),
                    ("Get.inputs:primPath", p_in.GetPrimPath().pathString),
                    ("Set.inputs:primPath", p_out.GetPrimPath().pathString),
                ],
            },
        )
        await omni.kit.app.get_app().next_update_async()

        event_counter = [0]

        def on_node_event(event):
            event_counter[0] += 1
            self.assertEqual(event.type, int(og.NodeEvent.ATTRIBUTE_TYPE_RESOLVE))

        sub = get_node.get_event_stream().create_subscription_to_pop(on_node_event)

        for _ in range(5):
            await omni.kit.app.get_app().next_update_async()
        self.assertLess(event_counter[0], 5)

        del sub

    # ----------------------------------------------------------------------
    async def test_loading_extended_attribute_values(self):
        """Test that saved extended attribute values are de-serialized correctly"""
        (result, error) = await ogts.load_test_file("TestExtendedAttributes.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        graph_path = "/World/PushGraph"
        controller = og.Controller()
        graph = controller.graph(graph_path)
        await controller.evaluate(graph)

        tests = [
            # name, og_type, input0, input1
            ("make_array", "bool", True, False, None),
            ("make_array_01", "double", 42, 43, None),
            ("make_array_02", "token", "FOO", "BAR", None),
            ("make_array_03", "double[2]", np.array((1.0, 2.0)), np.array((3.0, 4.0)), Gf.Vec2d),
            ("make_array_04", "int[2]", np.array((1, 2)), np.array((3, 4)), Gf.Vec2i),
            ("make_array_05", "half", 42.0, 43.0, None),
            (
                "make_array_06",
                "matrixd[3]",
                np.array((1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0)),
                np.array((1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0)),
                Gf.Matrix3d,
            ),
        ]

        # verify the values on the loaded nodes
        for name, og_type, expected_1, expected_2, _ in tests:
            make_array = controller.node(f"{graph_path}/{name}")
            input_0 = og.Controller.get(controller.attribute("inputs:input0", make_array))
            input_type = controller.attribute("inputs:input0", make_array).get_resolved_type().get_ogn_type_name()
            input_1 = og.Controller.get(controller.attribute("inputs:input1", make_array))
            value = og.Controller.get(controller.attribute("outputs:array", make_array))
            self.assertEqual(input_type, og_type)
            try:
                self.assertEqual(input_0, expected_1)
            except ValueError:
                self.assertListEqual(list(input_0), list(expected_1))
            try:
                self.assertEqual(input_1, expected_2)
            except ValueError:
                self.assertListEqual(list(input_1), list(expected_2))

            try:
                self.assertListEqual(list(value), [expected_1, expected_2])
            except ValueError:
                self.assertListEqual([list(v) for v in value], [list(expected_1), list(expected_2)])

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

        # Reset the arrayType to clear the saved resolved values
        for name, _, _, _, _ in tests:
            og.Controller.set(controller.attribute(f"{graph_path}/{name}.inputs:arrayType"), "auto")

        # Verify that all the inputs are now unresolved
        for name, _, _, _, _ in tests:
            make_array = controller.node(f"{graph_path}/{name}")
            attrib = controller.attribute("inputs:input0", make_array)
            tp = attrib.get_resolved_type()
            self.assertEqual(tp.base_type, og.BaseDataType.UNKNOWN)
            attrib = controller.attribute("inputs:input1", make_array)
            tp = attrib.get_resolved_type()
            self.assertEqual(tp.base_type, og.BaseDataType.UNKNOWN)

        # Verify that in the re-serialized stage, the attrType and attrValue are reset
        await controller.evaluate(graph)
        for name, _, _, _, _ in tests:
            self.assertIsNone(
                stage.GetAttributeAtPath(f"{graph_path}/{name}.inputs:input1").GetCustomDataByKey(
                    "omni:graph:attrValue"
                )
            )
            self.assertEqual(
                "Any",
                stage.GetAttributeAtPath(f"{graph_path}/{name}.inputs:input1").GetCustomDataByKey(
                    "omni:graph:attrType"
                ),
            )
            self.assertIsNone(
                stage.GetAttributeAtPath(f"{graph_path}/{name}.inputs:input1").GetCustomDataByKey(
                    "omni:graph:resolvedType"
                )
            )

        # Re-resolve and set different values
        for name, og_type, _, input_1, _ in tests:
            og.Controller.set(controller.attribute(f"{graph_path}/{name}.inputs:arrayType"), f"{og_type}[]")
            og.Controller.set(controller.attribute(f"{graph_path}/{name}.inputs:input1"), input_1)

        # Serialize stage again, verify the types and values
        await controller.evaluate(graph)
        for name, og_type, _, _, _ in tests:
            attr = stage.GetAttributeAtPath(f"{graph_path}/{name}.inputs:input1")
            self.assertIsNotNone(attr.GetCustomDataByKey("omni:graph:attrValue"), f"For {name}")
            self.assertIsNotNone(attr.GetCustomDataByKey("omni:graph:attrType"), f"For {name}")
            usd_value = stage.GetAttributeAtPath(f"{graph_path}/{name}.inputs:input1").GetCustomDataByKey(
                "omni:graph:resolvedType"
            )
            self.assertEqual(og_type, usd_value, f"For {name}")

    # ----------------------------------------------------------------------
    async def test_resolution_invalidates_existing_connection(self):
        """
        Test that when the type resolution makes an existing connection invalid, that we get an error
        """
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (build_str, find_prims), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("BuildString", "omni.graph.nodes.BuildString"),
                    ("FindPrims", "omni.graph.nodes.FindPrims"),
                ],
                keys.CREATE_PRIMS: [("/World/Cube", "Cube"), ("/World/Cone", "Cone")],
                keys.SET_VALUES: [("FindPrims.inputs:rootPrimPath", "/World")],
                keys.CONNECT: [("BuildString.outputs:value", "FindPrims.inputs:namePrefix")],
            },
        )

        await og.Controller.evaluate(graph)

        out = find_prims.get_attribute("outputs:prims")
        self.assertEqual(out.get(), [f"{self.TEST_GRAPH_PATH}", "/World/Cube", "/World/Cone"])

        # Setting inputs to string will change resolution on output value to a string as well. But string -> token
        # connections are not supported. So this should raise an error
        with ogts.ExpectedError():
            controller.edit(
                self.TEST_GRAPH_PATH,
                {keys.SET_VALUES: [("BuildString.inputs:a", {"type": "string", "value": "Cu"})]},
            )
            err_list = build_str.get_compute_messages(og.Severity.ERROR)
            self.assertEqual(len(err_list), 1)
            err_list[0].startswith("Type error")
            await controller.evaluate(graph)

        out = build_str.get_attribute("outputs:value")
        self.assertEqual(out.get_resolved_type(), og.Controller.attribute_type("string"))

    # ----------------------------------------------------------------------
    async def test_resolved_type_saved(self):
        """Test that a type resolution is saved when a value is never set"""
        # See OM-93596
        controller = og.Controller()
        keys = og.Controller.Keys
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Add", "omni.graph.nodes.Add"),
                    ("ConstantFloat", "omni.graph.nodes.ConstantFloat"),
                ],
                keys.SET_VALUES: [
                    ("Add.inputs:a", {"type": "float", "value": 1.0}),
                    ("ConstantFloat.inputs:value", 42.0),
                ],
            },
        )
        await controller.evaluate(graph)
        attr = stage.GetAttributeAtPath(f"{self.TEST_GRAPH_PATH}/Add.inputs:a")
        og_attr = controller.attribute(f"{self.TEST_GRAPH_PATH}/Add.inputs:a")
        self.assertIsNotNone(attr.GetCustomDataByKey("omni:graph:attrValue"))
        # Since the controller manually resolves the type above, verify the customData is there
        usd_value = attr.GetCustomDataByKey("omni:graph:resolvedType")
        self.assertEqual("float", usd_value)

        # manually resolve to a different type - verify it is saved
        og.cmds.ResolveAttrType(attr=og_attr, type_id="double")
        await controller.evaluate(graph)
        usd_value = attr.GetCustomDataByKey("omni:graph:resolvedType")
        self.assertEqual("double", usd_value)

        # manually unresolve - verify the customData is gone
        og.cmds.ResolveAttrType(attr=og_attr, type_id="unknown")
        await controller.evaluate(graph)
        usd_value = attr.GetCustomDataByKey("omni:graph:resolvedType")
        self.assertIsNone(usd_value)
        self.assertEqual(og_attr.get_resolved_type(), og.Type(og.BaseDataType.UNKNOWN, 1, 0))

        # Verify that:
        # 1. manually resolving
        # 2. connect
        # 3. disconnect
        # Restores the manual resolution customData AND literal value

        # FIXME: OM-94077
        if False:  # noqa PLW0125
            og.cmds.ResolveAttrType(attr=og_attr, type_id="float")
            og.Controller.set(attr, {"type": "float", "value": 2.0})
            await controller.evaluate(graph)
            self.assertEqual(attr.GetCustomDataByKey("omni:graph:attrValue"), 2.0)
            usd_value = attr.GetCustomDataByKey("omni:graph:resolvedType")
            self.assertEqual("float", usd_value)
            controller.edit(self.TEST_GRAPH_PATH, {keys.CONNECT: [("ConstantFloat.inputs:value", "Add.inputs:a")]})
            await controller.evaluate(graph)
            self.assertEqual(og_attr.get(), 42.0)
            controller.edit(
                self.TEST_GRAPH_PATH,
                {keys.DISCONNECT: [("ConstantFloat.inputs:value", "Add.inputs:a")]},
            )
            await controller.evaluate(graph)
            self.assertEqual(attr.GetCustomDataByKey("omni:graph:attrValue"), 2.0)
            usd_value = attr.GetCustomDataByKey("omni:graph:resolvedType")
            self.assertEqual("float", usd_value)
            self.assertEqual(og_attr.get(), 2.0)
            controller.edit(self.TEST_GRAPH_PATH, {keys.CONNECT: [("Add.outputs:sum", "ConstantFloat.inputs:value")]})
            await controller.evaluate(graph)
            og_sum_attr = controller.attribute(f"{self.TEST_GRAPH_PATH}/Add.outputs:sum")
            self.assertEqual(og_sum_attr.get_resolved_type(), og.Type(og.BaseDataType.FLOAT, 1, 0))
            controller.edit(
                self.TEST_GRAPH_PATH,
                {keys.DISCONNECT: [("Add.outputs:sum", "ConstantFloat.inputs:value")]},
            )
            await controller.evaluate(graph)
            self.assertEqual(og_sum_attr.get_resolved_type(), og.Type(og.BaseDataType.UNKNOWN, 1, 0))

    # ----------------------------------------------------------------------
    async def test_resolved_types_serialized_in_compounds(self):
        """Validates that resolved types are serialized in compound graphs (OM-99390)"""

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

        (_, _, _, nodes) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Add", "omni.graph.nodes.Add"),
                    (
                        "Compound",
                        {
                            keys.CREATE_NODES: [("CompoundAdd", "omni.graph.nodes.Add")],
                            keys.SET_VALUES: [("CompoundAdd.inputs:a", {"type": "float", "value": 2.0})],
                        },
                    ),
                ],
                keys.SET_VALUES: [
                    ("Add.inputs:a", {"type": "double", "value": 1.0}),
                ],
            },
        )

        add_path = nodes["Add"].get_prim_path()
        compound_add_path = nodes["CompoundAdd"].get_prim_path()

        # resolve and set attributes
        top_level_attr_a = nodes["Add"].get_attribute("inputs:a")
        top_level_attr_b = nodes["Add"].get_attribute("inputs:b")
        compound_attr_a = nodes["CompoundAdd"].get_attribute("inputs:a")
        compound_attr_b = nodes["CompoundAdd"].get_attribute("inputs:b")

        og.cmds.ResolveAttrType(attr=top_level_attr_b, type_id="double")
        og.cmds.ResolveAttrType(attr=compound_attr_b, type_id="float")
        controller.set(top_level_attr_b, 3.0)
        controller.set(compound_attr_b, 4.0)

        self.assertEqual(og.Controller.get(top_level_attr_a), 1.0)
        self.assertEqual(og.Controller.get(compound_attr_a), 2.0)
        self.assertEqual(og.Controller.get(top_level_attr_b), 3.0)
        self.assertEqual(og.Controller.get(compound_attr_b), 4.0)

        # save and reload the scene
        with tempfile.TemporaryDirectory() as test_directory:
            tmp_file_path = os.path.join(test_directory, "tmp_test_resolve_in_compounds.usda")
            (result, error, _) = await omni.usd.get_context().save_as_stage_async(str(tmp_file_path))
            self.assertTrue(result, error)

            (result, error) = await omni.usd.get_context().reopen_stage_async()
            self.assertTrue(result, error)

        add_node = og.Controller.node(add_path)
        compound_add_node = og.Controller.node(compound_add_path)

        top_level_attr_a = add_node.get_attribute("inputs:a")
        top_level_attr_b = add_node.get_attribute("inputs:b")
        compound_attr_a = compound_add_node.get_attribute("inputs:a")
        compound_attr_b = compound_add_node.get_attribute("inputs:b")

        self.assertEqual(og.Controller.get(top_level_attr_a), 1.0)
        self.assertEqual(og.Controller.get(compound_attr_a), 2.0)
        self.assertEqual(og.Controller.get(top_level_attr_b), 3.0)
        self.assertEqual(og.Controller.get(compound_attr_b), 4.0)
