"""Basic tests of the controller classes in omni.graph.core"""
import json
from typing import Any, Dict, List

import carb
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
import omni.usd
from omni.graph.core._impl.utils import _flatten_arguments, _Unspecified
from pxr import Sdf

_KEYS = og.Controller.Keys
"""Syntactic sugar to shorten the name required to access the keywords"""


# ==============================================================================================================
class TestControllers(ogts.OmniGraphTestCase):
    """Run a simple set of unit tests that exercise the main controller functionality"""

    # --------------------------------------------------------------------------------------------------------------
    async def test_arg_flattening(self):
        """Test the utility that flattens arguments list into a common dictionary"""
        test_data = [
            ([], [("a", 1)], [], {}, [1]),
            ([("a", _Unspecified)], [], [1], {}, [1]),
            ([("a", 1)], [], [], {}, [1]),
            ([("a", _Unspecified)], [], [], {"a": 2}, [2]),
            ([("a", _Unspecified)], [], [5], {}, [5]),
            ([("a", _Unspecified)], [("b", 7)], [5], {}, [5, 7]),
            ([("a", _Unspecified)], [("b", 7)], [], {"a": 3}, [3, 7]),
            ([("a", _Unspecified)], [("b", 7)], [5], {"b": 3}, [5, 3]),
            ([("a", _Unspecified)], [("b", _Unspecified)], [5], {"b": 3}, [5, 3]),
        ]

        for test_index, (mandatory, optional, args, kwargs, expected_results) in enumerate(test_data):
            flat_args = _flatten_arguments(mandatory, optional, args, kwargs)
            self.assertCountEqual(flat_args, expected_results, f"Test {test_index}")

        # Test cases where the arguments are invalid and exceptions are expected
        test_data_invalid = [
            ([("a", _Unspecified)], [], [], {}, [1]),  # Missing mandatory arg
            ([("a", _Unspecified)], [], [], {"b": 1}, []),  # Unknown kwarg
            ([("a", _Unspecified)], [], [1, 2], {}, [1]),  # Too many args
            ([("a", _Unspecified)], [], [1], {"a": 2}, [1]),  # Ambiguous arg value
            ([("a", _Unspecified)], [("b", 7)], [5], {"a": 3}, [5, 3]),  # Ambiguous arg value
        ]
        for mandatory, optional, args, kwargs, _expected_results in test_data_invalid:
            with self.assertRaises(og.OmniGraphError):
                _ = _flatten_arguments(mandatory, optional, args, kwargs)
                carb.log_error(f"Should have failed with {mandatory}, {optional}, {args}, {kwargs}")

    # --------------------------------------------------------------------------------------------------------------
    async def test_flattening_example(self):
        """Test that runs the code cited in the documentation for the _flatten_arguments function"""
        feet = 1
        yards = 3
        miles = 5280

        class Toady:
            def __init__(self):
                self.__height_default = 1
                self.__height_unit = yards
                self.jump = self.__jump_obj

            @classmethod
            def jump(cls, *args, **kwargs):  # noqa: PLE0202  Hiding it is the point
                return cls.__jump(cls, args=args, kwargs=kwargs)

            def __jump_obj(self, *args, **kwargs):
                return self.__jump(
                    self, how_high=self.__height_default, unit=self.__height_unit, args=args, kwargs=kwargs
                )

            @staticmethod
            def __jump(obj, how_high: int = None, unit: str = miles, args=List[Any], kwargs=Dict[str, Any]):
                (how_high, unit) = _flatten_arguments(
                    mandatory=[("how_high", how_high)], optional=[("unit", unit)], args=args, kwargs=kwargs
                )
                return how_high, unit

        # Legal calls to this function
        self.assertEqual((123, miles), Toady.jump(how_high=123))
        self.assertEqual((123, miles), Toady.jump(123))
        self.assertEqual((123, feet), Toady.jump(123, unit=feet))
        self.assertEqual((123, feet), Toady.jump(how_high=123, unit=feet))

        # The main difference in object-based calls is the use of the object members for defaults, also allowing it
        # to omit the "mandatory" members in the call.
        smithers = Toady()
        self.assertEqual((1, yards), smithers.jump())
        self.assertEqual((123, yards), smithers.jump(how_high=123))
        self.assertEqual((123, yards), smithers.jump(123))
        self.assertEqual((123, feet), smithers.jump(123, unit=feet))
        self.assertEqual((123, feet), smithers.jump(how_high=123, unit=feet))
        self.assertEqual((1, feet), smithers.jump(unit=feet))

    # --------------------------------------------------------------------------------------------------------------
    async def test_graph_construction(self):
        """Test the basic graph construction function of the og.Controller class"""
        controller = og.Controller()

        (graph, _, _, _) = controller.edit("/World/PushGraph")
        self.assertTrue(graph, "Created a new graph")
        await controller.evaluate(graph)

        graph_by_path = og.get_graph_by_path("/World/PushGraph")
        self.assertEqual(graph, graph_by_path, "Created graph lookup")

        self.assertEqual(graph.get_pipeline_stage(), og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION)
        self.assertEqual(graph.get_graph_backing_type(), og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED)

        (same_graph, _, _, _) = controller.edit(graph)
        self.assertEqual(graph, same_graph, "Edit of the same graph twice")
        await controller.evaluate(same_graph)

        index = 0
        all_graphs = [graph]
        for backing_type in og.GraphBackingType.__members__.values():
            if backing_type in [
                og.GraphBackingType.GRAPH_BACKING_TYPE_UNKNOWN,
                og.GraphBackingType.GRAPH_BACKING_TYPE_NONE,
                og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_WITH_HISTORY,
            ]:
                continue
            for pipeline_stage in og.GraphPipelineStage.__members__.values():
                if pipeline_stage in [og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_UNKNOWN]:
                    continue

                for evaluation_mode in [
                    og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_AUTOMATIC,
                    og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_STANDALONE,
                    og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_INSTANCED,
                ]:
                    error_msg = (
                        f"Created graph with backing {backing_type}, pipeline {pipeline_stage} "
                        "and evaluation mode {evaluation_mode}"
                    )
                    (new_graph, _, _, _) = controller.edit(
                        {
                            "graph_path": f"/World/TestGraph{index}",
                            "fc_backing_type": backing_type,
                            "pipeline_stage": pipeline_stage,
                            "evaluation_mode": evaluation_mode,
                        }
                    )

                    self.assertTrue(new_graph, error_msg)
                    self.assertEqual(new_graph.get_pipeline_stage(), pipeline_stage, error_msg)
                    self.assertEqual(new_graph.get_graph_backing_type(), backing_type, error_msg)
                    self.assertEqual(new_graph.evaluation_mode, evaluation_mode, error_msg)
                    self.assertTrue(
                        any(
                            new_graph == pipeline_graph
                            for pipeline_graph in og.get_graphs_in_pipeline_stage(pipeline_stage)
                        )
                    )
                    index += 1
                    all_graphs.append(new_graph)
                    await controller.evaluate(all_graphs)

    # ----------------------------------------------------------------------
    async def test_graph_population(self):
        """Test node creation and deletion via the og.Controller class"""
        controller = og.Controller()

        simple_node_type = "omni.graph.tutorials.SimpleData"

        # Create a couple of random nodes
        (graph, nodes_created, _, _) = controller.edit(
            "/World/PushGraph",
            {
                _KEYS.CREATE_NODES: [
                    ("Simple1", simple_node_type),
                    ("Simple2", simple_node_type),
                    ("Simple3", simple_node_type),
                ]
            },
        )
        for node in nodes_created:
            self.assertTrue(node.is_valid(), "Created node is valid")
            self.assertEqual(node.get_node_type().get_node_type(), simple_node_type)

        self.assertEqual(len(nodes_created), 3, "Correct number of nodes created and returned")

        nodes_matched = 0
        for node in graph.get_nodes():
            for expected_node in nodes_created:
                if node == expected_node:
                    nodes_matched += 1
                    break
        self.assertEqual(nodes_matched, 3, "Found all created nodes in the return values")

        # Delete one node, add two nodes, then delete an earlier one.
        # Uses a static controller instead of the existing one for some operations to confirm that also works
        controller.edit(graph, {_KEYS.DELETE_NODES: "Simple1"})
        (static_graph, static_nodes, _, static_path_node_map) = og.Controller.edit(
            graph,
            {
                _KEYS.CREATE_NODES: [
                    ("Simple1", simple_node_type),
                    ("Simple4", simple_node_type),
                ]
            },
        )
        self.assertEqual(graph, static_graph, "Static controller uses same graph")
        self.assertEqual(len(static_nodes), 2, "Static node creation")
        self.assertCountEqual(["Simple1", "Simple4"], list(static_path_node_map.keys()), "Static node path map")
        (_, _, _, new_path_node_map) = controller.edit(graph, {_KEYS.DELETE_NODES: "Simple2"})
        self.assertCountEqual(
            ["Simple3"], list(new_path_node_map.keys()), "Path map is only aware of object-based changes"
        )

        # Delete a node then immediately create one with the same name but a different type, using the non-list forms
        (_, retyped_nodes, _, retyped_path_node_map) = controller.edit(
            graph,
            {
                _KEYS.DELETE_NODES: "Simple1",
                _KEYS.CREATE_NODES: ("Simple1", "omni.tutorials.TupleData"),
            },
        )
        self.assertCountEqual(
            ["Simple1", "Simple3"],
            list(retyped_path_node_map.keys()),
            "Path map deleted and added a node of the same name",
        )
        self.assertEqual(
            "omni.tutorials.TupleData",
            controller.node_type(retyped_nodes[0]).get_node_type(),
            "Same node name, different type",
        )

        # Attempt to add a node with a bad type
        with self.assertRaises(og.OmniGraphError):
            og.Controller.edit(graph, {_KEYS.CREATE_NODES: [("NoSuchNode", "omni.no.such.type")]})

        # Attempt to delete a non-existing node (using a static Controller object instead of the existing one)
        with self.assertRaises(og.OmniGraphError):
            og.Controller.edit(graph, {_KEYS.DELETE_NODES: ["/This/Node/Does/Not/Exist"]})

    # --------------------------------------------------------------------------------------------------------------
    async def test_prim_construction(self):
        """Test the controller's ability to create USD prims. This is just testing basic functionality. A separate
        test will exercise creation of prims with all of the available attribute types.

        Testing Type Matrix:
            prim_path: str, Sdf.Path, Invalid
            attribute_id: dict of
                KEY: str, Invalid
                VALUE: 2-tuple of
                    VALUE[0]: str(ogn type), str(sdf type), og.Type, Invalid
                    VALUE[1]: Matching Value Type, Unmatching Value Type, Invalid
            prim_type: str, None
        """
        controller = og.Controller()
        (graph, _, prims, _,) = controller.edit(
            "/PrimGraph",
            {
                controller.Keys.CREATE_PRIMS: [
                    # (str, {str: (ogn_str, value)})
                    ("PrimFloat2", {"float2Value": ("float[2]", [1.0, 2.0])}),
                    # (str, {Sdf.Path: (sdf_str, value)})
                    (Sdf.Path("PrimFloat3"), {"float3Value": ("float3", [1.0, 2.0, 3.0])}),
                    # (str, {str: (og.Type, value)})
                    ("PrimFloat4", {"float4Value": (og.Type(og.BaseDataType.FLOAT, 4, 0), [1.0, 2.0, 3.0, 4.0])}),
                    # str
                    "PrimWithNoAttributes",
                ]
            },
        )
        # Verify prims are created with the right attributes
        self.assertEqual(len(prims), 4)
        self.assertTrue(all(prim.IsValid() for prim in prims))

        float2_attr = prims[0].GetAttribute("float2Value")
        self.assertTrue(float2_attr.IsValid())
        self.assertEqual([1.0, 2.0], float2_attr.Get())

        float3_attr = prims[1].GetAttribute("float3Value")
        self.assertTrue(float3_attr.IsValid())
        self.assertEqual([1.0, 2.0, 3.0], float3_attr.Get())

        float4_attr = prims[2].GetAttribute("float4Value")
        self.assertTrue(float4_attr.IsValid())
        self.assertEqual([1.0, 2.0, 3.0, 4.0], float4_attr.Get())

        # Test that non-string prim path fails
        with self.assertRaises(og.OmniGraphError):
            og.Controller.edit(graph, {controller.Keys.CREATE_PRIMS: [prims[0]]})

        # Test that attempt to create prim in already existing location fails, with both absolute and relative paths
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {controller.Keys.CREATE_PRIMS: [("PrimFloat2", "Cube")]})
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {controller.Keys.CREATE_PRIMS: [("/PrimFloat2", "Cube", {})]})

        # Test that prims are forbidden from being created inside an OmniGraph
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {controller.Keys.CREATE_PRIMS: [("/PrimGraph/IllegalLocation", {})]})

        # Test that non-string attribute name fails
        with self.assertRaises(og.OmniGraphError):
            controller.edit(
                graph,
                {
                    controller.Keys.CREATE_PRIMS: [
                        ("PrimWithBadAttribute", {None: None}),
                    ]
                },
            )

        # Test that invalid attribute data type fails
        with self.assertRaises(og.OmniGraphError):
            controller.edit(
                graph,
                {
                    controller.Keys.CREATE_PRIMS: [
                        ("PrimWithBadAttributeData", {"complex": [1.0, 2.0]}),
                    ]
                },
            )

        # Test that non-matching attribute data type fails
        with self.assertRaises(og.OmniGraphError):
            controller.edit(
                graph,
                {
                    controller.Keys.CREATE_PRIMS: [
                        ("PrimWithBadAttributeData", {"float": None}),
                    ]
                },
            )

        # Test that invalid attribute data fails
        with self.assertRaises(og.OmniGraphError):
            controller.edit(
                graph,
                {
                    controller.Keys.CREATE_PRIMS: [
                        ("PrimWithBadAttributeData", {"float": None}),
                    ]
                },
            )

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

    # --------------------------------------------------------------------------------------------------------------
    async def test_connections(self):
        """Test the various connection-related functions on the graph controller."""

        # Use the main controller for testing as it derives from the graph controller
        controller = og.Controller()

        # Source --> Sink    --> InOut
        #       \    \      /
        #        \    -----/
        #         --> FanOut
        # Command 1
        (graph, nodes_created, _, _) = controller.edit(
            "/TestGraph",
            {
                controller.Keys.CREATE_NODES: [
                    ("Source", "omni.graph.tutorials.SimpleData"),
                    ("Sink", "omni.graph.tutorials.SimpleData"),
                    ("FanOut", "omni.graph.tutorials.SimpleData"),
                    ("InOut", "omni.graph.tutorials.SimpleData"),
                ],
                controller.Keys.CONNECT: [
                    ("Source.outputs:a_bool", "Sink.inputs:a_bool"),
                    (("outputs:a_bool", "Source"), ("inputs:a_bool", "FanOut")),
                    ("Sink.inputs:a_bool", "InOut.inputs:a_bool"),
                ],
            },
        )
        (source_node, sink_node, fanout_node, inout_node) = nodes_created

        def _get_nodes():
            return (
                controller.node(("Source", graph)),
                controller.node(("Sink", graph)),
                controller.node(("FanOut", graph)),
                controller.node(("InOut", graph)),
            )

        def _get_attributes():
            return (
                controller.attribute(("outputs:a_bool", source_node)),
                controller.attribute("inputs:a_bool", sink_node),
                controller.attribute(("inputs:a_bool", fanout_node)),
                controller.attribute("inputs:a_bool", inout_node),
            )

        source_output, sink_input, fanout_input, inout_input = _get_attributes()

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

        # Command 2
        controller.disconnect_all(("outputs:a_bool", source_node))
        self.__confirm_connection_counts(source_output, 0, 0)
        self.__confirm_connection_counts(sink_input, 0, 1)
        self.__confirm_connection_counts(fanout_input, 0, 0)
        self.__confirm_connection_counts(inout_input, 1, 0)

        # Command 3
        controller.disconnect_all(sink_input)
        self.__confirm_connection_counts(source_output, 0, 0)
        self.__confirm_connection_counts(sink_input, 0, 0)
        self.__confirm_connection_counts(fanout_input, 0, 0)
        self.__confirm_connection_counts(inout_input, 0, 0)

        # Command 4
        controller.edit(graph, {controller.Keys.DELETE_NODES: [source_node, sink_node, fanout_node, inout_node]})
        omni.kit.undo.undo()  # Command 4
        omni.kit.undo.undo()  # Command 3

        # May have lost the objects through the undo process so get them again
        (source_node, sink_node, fanout_node, inout_node) = _get_nodes()
        source_output, sink_input, fanout_input, inout_input = _get_attributes()

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

        # Go back and forth in the undo queue to get back to having all objects and connections present
        omni.kit.undo.redo()  # Command 2
        omni.kit.undo.redo()  # Command 3
        omni.kit.undo.redo()  # Command 4
        omni.kit.undo.undo()  # Command 4
        omni.kit.undo.undo()  # Command 3
        omni.kit.undo.undo()  # Command 2

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

        # Direct calls are okay when the attribute is fully specified
        # Command 5
        og.Controller.disconnect_all(sink_input)
        self.__confirm_connection_counts(source_output, 0, 1)
        self.__confirm_connection_counts(sink_input, 0, 0)
        self.__confirm_connection_counts(fanout_input, 1, 0)
        self.__confirm_connection_counts(inout_input, 0, 0)

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

    # --------------------------------------------------------------------------------------------------------------
    async def test_prim_exposure(self):
        """Test the controller's ability to expose USD prims to OmniGraph.
        def expose_prim(cls, exposure_type: PrimExposureType, prim_id: Prim_t, node_path_id: NewNode_t)
            -> Union[og.Node, List[og.Node]]:

        The test proceeds by first creating a few prims and then exposing them to OmniGraph, confirming that the
        OmniGraph nodes do in fact get the correct values from their underlying prims.
        """
        controller = og.Controller()
        exposure_types = og.GraphController.PrimExposureType
        (graph, _, _, path_to_object_map,) = controller.edit(
            "/PrimGraph",
            {
                _KEYS.CREATE_PRIMS: [
                    ("PrimFloat2", {"float2Value": ("float[2]", [1.0, 2.0])}),
                    ("Cube", "Cube"),
                    ("PrimVelocity", {"velocity": ("float[3]", [1.0, 2.0, 3.0])}, "Velocity"),
                    ("Prim1"),
                    ("Prim2"),
                    ("Prim3"),
                    ("PrimPosition", {"position": ("pointf[3]", [1.0, 2.0, 3.0])}, "Position"),
                ]
            },
        )

        # Test exposure of a prim as a bundle
        (_, exposed_nodes, _, _) = og.Controller.edit(
            graph,
            {
                _KEYS.EXPOSE_PRIMS: (exposure_types.AS_BUNDLE, "PrimFloat2", "ExposedAsBundle"),
            },
            path_to_object_map,
        )
        await controller.evaluate(graph)
        self.assertEqual(len(exposed_nodes), 1, "Exposed a simple prim by bundle")
        self.assertTrue("ExposedAsBundle" in path_to_object_map)
        self.assertEqual(exposed_nodes[0].get_prim_path(), "/PrimGraph/ExposedAsBundle")
        self.assertEqual(exposed_nodes[0].get_type_name(), "omni.graph.nodes.ExtractPrim")
        (_, inspector_nodes, _, _) = controller.edit(
            graph,
            {
                _KEYS.CREATE_NODES: ("Inspector", "omni.graph.nodes.BundleInspector"),
                _KEYS.CONNECT: ("ExposedAsBundle.outputs_primBundle", "Inspector.inputs:bundle"),
            },
        )
        await controller.evaluate(graph)
        # float2Value is in the exposed bundle
        inspector_node = inspector_nodes[0]
        bundle_names = og.Controller.get(controller.attribute("outputs:names", inspector_node))
        index = bundle_names.index("float2Value")
        self.assertTrue(index >= 0)

        # The value of that attribute is [1.0, 2.0]
        try:
            bundle_values = og.Controller.get(controller.attribute("outputs:values", inspector_node))
            # This weird thing is necessary because the tuple values are not JSON-compatible
            tuple_to_list = json.loads(bundle_values[index].replace("(", "[").replace(")", "]"))
            self.assertCountEqual(tuple_to_list, [1.0, 2.0])
        except AssertionError as error:
            carb.log_warn(f"ReadBundle node can't get the bundle to the inspector - {error}")

        # Test exposure of a prim as attributes
        (_, exposed_nodes, _, _) = og.Controller.edit(
            graph,
            {
                _KEYS.EXPOSE_PRIMS: (exposure_types.AS_ATTRIBUTES, "PrimVelocity", "ExposedAsAttribute"),
            },
            path_to_object_map,
        )
        await controller.evaluate(graph)
        self.assertEqual(len(exposed_nodes), 1, "Exposed a simple prim by attribute")
        self.assertTrue("ExposedAsAttribute" in path_to_object_map)
        self.assertEqual(exposed_nodes[0].get_prim_path(), "/PrimGraph/ExposedAsAttribute")
        self.assertEqual(exposed_nodes[0].get_type_name(), "omni.graph.nodes.ReadPrimAttributes")

        # in order to read data, the attribute needs to be connected, and an eval must be done
        (graph, _, _, _) = og.Controller.edit(
            graph,
            {
                _KEYS.CREATE_NODES: [
                    ("Const", "omni.graph.nodes.ConstantFloat3"),
                ],
            },
        )

        # we need to evaluate for dynamic attributes to appear on ExtractPrim2
        await controller.evaluate(graph)

        # after attributes appeared on ExtractBundle2 we can connect them
        (graph, _, _, _) = og.Controller.edit(
            graph,
            {
                _KEYS.CONNECT: ("/PrimGraph/ExposedAsAttribute.outputs:velocity", "/PrimGraph/Const.inputs:value"),
            },
        )
        await controller.evaluate(graph)

        self.assertCountEqual(
            [1.0, 2.0, 3.0], og.Controller.get(controller.attribute(("outputs:velocity", exposed_nodes[0])))
        )

        # Test exposure of a list of prims
        (_, exposed_nodes, _, _) = og.Controller.edit(
            graph,
            {
                _KEYS.EXPOSE_PRIMS: [
                    (exposure_types.AS_ATTRIBUTES, "Prim1", "Exposed1"),
                    (exposure_types.AS_ATTRIBUTES, "Prim2", "Exposed2"),
                    (exposure_types.AS_ATTRIBUTES, "Prim3", "Exposed3"),
                ]
            },
            path_to_object_map,
        )
        await controller.evaluate(graph)
        self.assertEqual(len(exposed_nodes), 3, "Exposed a list of prims by attribute")
        for index in range(3):
            self.assertTrue(f"Exposed{index + 1}" in path_to_object_map)
            self.assertEqual(exposed_nodes[index].get_prim_path(), f"/PrimGraph/Exposed{index + 1}")
            self.assertEqual(exposed_nodes[index].get_type_name(), "omni.graph.nodes.ReadPrimAttributes")

        # Test exposure as write
        (_, exposed_nodes, _, _) = controller.edit(
            graph, {_KEYS.EXPOSE_PRIMS: (exposure_types.AS_WRITABLE, "PrimPosition", "ExposedWritable")}
        )
        await controller.evaluate(graph)
        write_node = exposed_nodes[0]
        self.assertEqual(len(exposed_nodes), 1, "Exposed a simple writable prim")
        self.assertTrue("ExposedWritable" in path_to_object_map)
        self.assertEqual(write_node.get_prim_path(), "/PrimGraph/ExposedWritable")
        self.assertEqual(write_node.get_type_name(), "omni.graph.nodes.WritePrim")
        # Check we have the expected dynamic attrib on WritePrim
        attribs = write_node.get_attributes()
        found_size_attrib = False
        for attrib in attribs:
            if attrib.get_name() == "inputs:position" and attrib.get_resolved_type().get_ogn_type_name() == "pointf[3]":
                found_size_attrib = True
        self.assertTrue(found_size_attrib)

        # Test invalid attempt to expose non-existent prim
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {_KEYS.EXPOSE_PRIMS: (exposure_types.AS_ATTRIBUTES, "NotAPrim", "NotANode")})

        # Test invalid attempt to expose prim on node path that's not in a graph
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {_KEYS.EXPOSE_PRIMS: (exposure_types.AS_ATTRIBUTES, "Cube", "/NotANode")})

        # Test invalid attempt to expose prim on an already existing node
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {_KEYS.EXPOSE_PRIMS: (exposure_types.AS_ATTRIBUTES, "Cube", "Exposed1")})

    # --------------------------------------------------------------------------------------------------------------
    async def test_set_values(self):
        """Test the controller's ability to set values on nodes.
        def set(AttributeValues_t)
        """
        controller = og.Controller()
        (graph, (node1, _, _, add_node), _, _) = controller.edit(
            "/TestGraph",
            {
                _KEYS.CREATE_NODES: [
                    ("Node1", "omni.graph.tutorials.SimpleData"),
                    ("Node2", "omni.graph.tutorials.SimpleData"),
                    ("Node3", "omni.graph.tutorials.SimpleData"),
                    ("AddNode", "omni.graph.nodes.Add"),
                ],
            },
        )

        # Test setting of a regular attribute value
        controller.edit(graph, {_KEYS.SET_VALUES: [("Node1.inputs:a_float", 17.0)]})
        await controller.evaluate(graph)
        # Exercise the controller version that takes an attribute (passing it up to the DataView)
        data_controller = og.Controller(attribute="/TestGraph/Node1.outputs:a_float")
        self.assertEqual(18.0, data_controller.get())

        # Test setting of an extended attribute value with a resolution type
        controller.edit(
            graph,
            {
                _KEYS.SET_VALUES: [
                    (("inputs:a", add_node), og.TypedValue(17.0, "float")),
                    ("AddNode.inputs:b", og.TypedValue(51.0, "float")),
                ]
            },
        )
        await controller.evaluate(graph)
        # Until type resolution works when values are set as well as connections are made this cannot check the output
        self.assertEqual(17.0, og.DataView.get(og.ObjectLookup.attribute(("inputs:a", add_node))))

        # Test setting of multiple attributes in one shot
        controller.edit(
            graph,
            {
                _KEYS.SET_VALUES: [
                    ("Node1.inputs:a_float", 42.0),
                    ("Node1.inputs:a_double", 23.0),
                ]
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(43.0, og.DataView.get(og.ObjectLookup.attribute(("outputs:a_float", node1))))
        self.assertEqual(24.0, og.Controller.get("/TestGraph/Node1.outputs:a_double"))

        # Test invalid attempt to set non-existent attribute
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {_KEYS.SET_VALUES: ("Node1.inputs:a_not_there", 5.0)})

        # Test invalid attempt to set extended attribute with non-existent type
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {_KEYS.SET_VALUES: (("inputs:a", add_node), 17.0, "not_a_type")})

    # --------------------------------------------------------------------------------------------------------------
    async def test_variables(self):
        """Test variable methods via the og.Controller class

        def get_variable_default_value(cls, variable_id: Variable_t):
        def set_variable_default_value(cls, variable_id: Variable_t, value):

        """
        controller = og.Controller()
        (graph, _, _, _) = controller.edit(
            "/TestGraph",
            {
                _KEYS.CREATE_VARIABLES: [
                    ("Variable_1", og.Type(og.BaseDataType.FLOAT)),
                    ("Variable_2", og.Type(og.BaseDataType.INT)),
                ],
            },
        )

        variable_1_id = graph.find_variable("Variable_1")
        og.Controller.set_variable_default_value(variable_1_id, 5)
        self.assertEqual(og.Controller.get_variable_default_value(variable_1_id), 5)

        variable_2_id = (graph, "Variable_2")
        og.Controller.set_variable_default_value(variable_2_id, 6)
        self.assertEqual(og.Controller.get_variable_default_value(variable_2_id), 6)

        with self.assertRaises(og.OmniGraphError):
            og.Controller.set_variable_default_value(variable_1_id, "Not a float")
        with self.assertRaises(og.OmniGraphError):
            og.Controller.get_variable_default_value((graph, "InvalidVariable"))
        with self.assertRaises(og.OmniGraphError):
            og.Controller.set_variable_default_value((graph, "InvalidVariable"), 5)
        with self.assertRaises(og.OmniGraphError):
            og.Controller.get_variable_default_value(None)

    # --------------------------------------------------------------------------------------------------------------
    async def test_compound_subgraph_graph_population(self):
        """Test graph population when using compound subgraphs"""

        controller = og.Controller()

        simple_node_type = "omni.graph.tutorials.SimpleData"

        (graph, nodes, _, _path_node_map) = controller.edit(
            "/TestGraph",
            {
                _KEYS.CREATE_NODES: [
                    (
                        "Compound_1",
                        {_KEYS.CREATE_NODES: [("Compound_1_1", simple_node_type), ("Compound_1_2", simple_node_type)]},
                    ),
                    (
                        "Compound_2",
                        {
                            _KEYS.CREATE_NODES: [
                                (
                                    "Compound_2_1",
                                    {
                                        _KEYS.CREATE_NODES: [
                                            ("Compound_2_1_1", simple_node_type),
                                            ("Compound_2_1_2", simple_node_type),
                                        ]
                                    },
                                )
                            ]
                        },
                    ),
                ]
            },
        )

        # test graph level
        self.assertTrue(graph.is_valid())
        self.assertEqual(2, len(nodes), "Expected 2 nodes create in graph")
        self.assertTrue(nodes[0].is_compound_node())
        self.assertTrue(nodes[1].is_compound_node())

        # test sublevels are constructed correctly
        compound_1_graph = nodes[0].get_compound_graph_instance()
        self.assertTrue(compound_1_graph.is_valid())
        compound_1_graph_nodes = compound_1_graph.get_nodes()
        self.assertEqual(2, len(compound_1_graph_nodes))
        self.assertEqual(compound_1_graph_nodes[0].get_node_type().get_node_type(), simple_node_type)
        self.assertEqual(compound_1_graph_nodes[1].get_node_type().get_node_type(), simple_node_type)

        compound_2_graph = nodes[1].get_compound_graph_instance()
        self.assertTrue(compound_2_graph.is_valid())
        compound_2_graph_nodes = compound_2_graph.get_nodes()
        self.assertEqual(1, len(compound_2_graph_nodes))
        compound_2_1_graph = compound_2_graph_nodes[0].get_compound_graph_instance()
        self.assertTrue(compound_2_1_graph.is_valid())
        compound_2_1_graph_nodes = compound_2_1_graph.get_nodes()
        self.assertEquals(2, len(compound_2_1_graph_nodes))
        self.assertEqual(compound_2_1_graph_nodes[0].get_node_type().get_node_type(), simple_node_type)
        self.assertEqual(compound_2_1_graph_nodes[1].get_node_type().get_node_type(), simple_node_type)

        # try to create a node with the same name in the subgraph. All node names must be unique
        with self.assertRaises(og.OmniGraphError):
            controller.edit("/TestGraph", {_KEYS.CREATE_NODES: ("Compound_1_1", simple_node_type)})

        # node with the same name on a subgraph as the parent
        with self.assertRaises(og.OmniGraphError):
            controller.edit(
                nodes[0].get_compound_graph_instance(), {_KEYS.CREATE_NODES: ("Compound_1", simple_node_type)}
            )

        # allowed with a new controller though
        og.Controller.edit(
            nodes[0].get_compound_graph_instance(), {_KEYS.CREATE_NODES: ("Compound_1", simple_node_type)}
        )
        self.assertTrue(compound_1_graph.get_node(f"{compound_1_graph.get_path_to_graph()}/Compound_1").is_valid())

        # create another graph that uses the class lookup (og.Controller) instead of an object
        (_graph2, nodes2, _, path_node_map2) = og.Controller.edit(
            "/TestGraph2",
            {
                _KEYS.CREATE_NODES: [
                    ("Compound_1", {_KEYS.CREATE_NODES: ("Compound_1_1", simple_node_type)}),
                ]
            },
        )
        self.assertTrue(nodes2[0].is_valid())
        self.assertTrue(nodes2[0].is_compound_node())
        self.assertTrue(path_node_map2["Compound_1_1"].is_valid())
        self.assertFalse(path_node_map2["Compound_1_1"].is_compound_node())

    # --------------------------------------------------------------------------------------------------------------
    async def test_attribute_promotion(self):
        """Tests the functionality of promoting attributes using the Controller._KEYS.PROMOTE_ATTRIBUTES command
        and the NodeController.promote_attribute() method"""

        controller = og.Controller(update_usd=True, undoable=True)

        simple_node_type = "omni.graph.tutorials.SimpleData"

        # graph with a subgraphs
        (graph, (compound, _non_compound), _, path_node_map) = controller.edit(
            "/TestGraph",
            {
                _KEYS.CREATE_NODES: [
                    (
                        "Compound_1",
                        {
                            _KEYS.CREATE_NODES: [
                                ("Compound_1_1", simple_node_type),
                                ("Compound_1_2", simple_node_type),
                                (
                                    "Compound_1_3",
                                    {
                                        _KEYS.CREATE_NODES: [
                                            ("Compound_1_3_1", simple_node_type),
                                        ]
                                    },
                                ),
                            ]
                        },
                    ),
                    ("NonCompound", simple_node_type),
                ]
            },
        )

        # cannot promote on the top-level graph
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("NonCompound.inputs:a_int", "inputs:a_int")})

        subgraph = compound.get_compound_graph_instance()

        # promote basic inputs, using both the graph and subgraph
        controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_1.inputs:a_int", "inputs:a_int")})
        controller.edit(subgraph, {_KEYS.PROMOTE_ATTRIBUTES: [("Compound_1_2.outputs:a_int", "outputs:a_int")]})
        self.assertTrue(compound.get_attribute_exists("inputs:a_int"))
        self.assertTrue(compound.get_attribute_exists("outputs:a_int"))

        # inputs already connected should fail.
        with ogts.ExpectedError():
            with self.assertRaises(og.OmniGraphError):
                controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_1.inputs:a_int", "inputs:a_int_2")})

        # output already connected should pass
        controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_2.outputs:a_int", "outputs:a_int_2")})
        self.assertTrue(compound.get_attribute_exists("outputs:a_int_2"))

        # input already exists
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_2.inputs:a_int", "inputs:a_int")})

        # output already exists
        with self.assertRaises(og.OmniGraphError):
            controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_1.outputs:a_int", "outputs:a_int")})

        # input can be promoted as output
        controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_1.inputs:a_bool", "outputs:a_bool")})

        # output cannot be promoted as input
        controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_1.outputs:a_bool", "inputs:a_bool_2")})
        self.assertTrue(compound.get_attribute_exists("outputs:inputs:a_bool_2"))

        # promoted name without prefix, the prefix is automatically added
        controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_1.inputs:a_float", "a_float")})
        controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_1.outputs:a_float", "a_float")})
        self.assertTrue(compound.get_attribute_exists("inputs:a_float"))
        self.assertTrue(compound.get_attribute_exists("outputs:a_float"))

        # validate output only inputs promote as outputs
        controller.edit(subgraph, {_KEYS.CREATE_NODES: ("Compound_1_4", "omni.graph.nodes.ConstantDouble")})
        controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_4.inputs:value", "constant")})
        self.assertTrue(compound.get_attribute_exists("outputs:constant"))

        # similar tests but using the NodeController API directly
        compound_1_1_node = path_node_map["Compound_1_1"]
        compound_1_2_node = path_node_map["Compound_1_2"]
        og.NodeController.promote_attribute(og.Controller.attribute("inputs:a_half", compound_1_1_node), "a_half")
        og.NodeController.promote_attribute(og.Controller.attribute("outputs:a_half", compound_1_2_node), "a_half")
        self.assertTrue(compound.get_attribute_exists("inputs:a_half"))
        self.assertTrue(compound.get_attribute_exists("outputs:a_half"))

        # already exists
        with self.assertRaises(og.OmniGraphError):
            og.NodeController.promote_attribute(og.Controller.attribute("inputs:a_double", compound_1_1_node), "a_half")
        with self.assertRaises(og.OmniGraphError):
            og.NodeController.promote_attribute(
                og.Controller.attribute("outputs:a_double", compound_1_1_node), "a_half"
            )

        # promote attributes on a nested recursion level
        controller.edit(graph, {_KEYS.PROMOTE_ATTRIBUTES: ("Compound_1_3_1.inputs:a_int", "inputs:value")})
        self.assertTrue(path_node_map["Compound_1_3"].get_attribute_exists("inputs:value"))

    # ----------------------------------------------------------------------
    async def test_attribute_creation(self):
        """Test attribute creation via the og.Controller class"""
        controller = og.Controller(undoable=False)
        (graph, (node,), _, _) = controller.edit(
            "/TestGraph",
            {
                _KEYS.CREATE_NODES: ("TestNode", "omni.graph.test.TestDynamicAttributeRawData"),
                _KEYS.CREATE_ATTRIBUTES: [
                    ("TestNode.inputs:a_float", "float"),
                ],
            },
        )

        self.assertTrue(node.is_valid())
        a_float = node.get_attribute("inputs:a_float")
        self.assertTrue(a_float.is_valid())

        base_attribute_count = len(node.get_attributes())

        # A scattered assortment of spec combinations that covers all of the types accepted by the command
        test_configs = [
            ("a_double", "/TestGraph/TestNode.inputs:a_double", "double"),
            ("a_any", Sdf.Path("/TestGraph/TestNode.inputs:a_any"), "any"),
            ("a_real", ("inputs:a_real", node), ["float", "double"]),
            ("a_numbers", (("a_numbers", og.AttributePortType.INPUT), node), ["numerics"]),
            ("a_bundle", ("inputs:a_bundle", "TestNode", graph), "bundle"),
            ("a_mesh", (("a_mesh", og.AttributePortType.INPUT), "TestNode", graph), "pointd[3][]"),
        ]

        for attr_name, attr_spec, attr_type in test_configs:
            controller.edit(graph, {_KEYS.CREATE_ATTRIBUTES: (attr_spec, attr_type)})
            attr = node.get_attribute(f"inputs:{attr_name}")
            self.assertTrue(attr.is_valid())

        self.assertEqual(base_attribute_count + len(test_configs), len(node.get_attributes()))


# AttributeTypeSpec_t = str | og.Type | list[str]
# """Typing that identifies a regular or extended type definition
# 1. "any"
# 2. og.Type(og.BaseDataType.FLOAT, 3, 1, og.AttributeRole.VECTOR)
# 3. ["float", "integral_scalers"]
# """

# NewAttribute_t = Path_t | tuple[AttributeName_t, Node_t] | tuple[AttributeName_t, str, Graph_t]
# """Typing for the information required to uniquely identify an attribute that does not yet exist
# 1a. "/MyGraph/MyNode/inputs:new_attr"
# 1b. Sdf.Path("/MyGraph/MyNode/inputs:new_attr")
# 2a. ("inputs:new_attr", og.Controller.node("/MyGraph/MyNode"))
# 2b. (("new_attr", og.AttributePortType.INPUT), og.Controller.node("/MyGraph/MyNode"))
# 3a. ("inputs:new_attr", "MyNode", og.Controller.graph("/MyGraph"))
# 3b. (("new_attr", og.AttributePortType.INPUT), "MyNode", og.Controller.graph("/MyGraph"))
# """
