"""Tests exercising every supported OGN data type"""
from math import isnan
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Set

import carb
import numpy as np
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.graph.nodes.tests as ognts
import omni.graph.tools.ogn as ogn
import omni.kit.test
import omni.usd

# Fast lookup of matrix dimensions, to avoid sqrt call
_MATRIX_DIMENSIONS = {4: 2, 9: 3, 16: 4}


# ======================================================================
class TestOmniGraphAllTypes(ogts.OmniGraphTestCase):
    """Run a simple unit test that exercises all of the data types supported by OmniGraph.
    These tests do not exercise any, union, or bundle types, which already have plenty of their own tests.
    """

    # These are not part of the generated attribute pattern so they will be ignored. They could be tested separately
    # but that won't add any value to the tests.
    UNCHECKED_ATTRIBUTES = [
        "node:type",
        "node:typeVersion",
        "NodeWithAllTypes",
        "inputs:doNotCompute",
        "inputs:a_target",
        "outputs:a_target",
        "state:a_target",
        "inputs:a_bundle",
        "outputs_a_bundle",
        "state_a_bundle",
        "state:a_stringEmpty",
        "state:a_firstEvaluation",
    ]

    TEST_GRAPH_PATH = "/World/TestGraph"

    # ----------------------------------------------------------------------
    def __attribute_type_to_name(self, attribute_type: og.Type) -> str:
        """Converts an attribute type into the canonical attribute name used by the test nodes.
        The rules are:
            - prefix of a_
            - followed by name of attribute base type
            - followed by optional _N if the component count N is > 1
            - followed by optional '_array' if the type is an array

        Args:
            type: OGN attribute type to deconstruct

        Returns:
            Canonical attribute name for the attribute with the given type
        """
        attribute_name = f"a_{og.Type(attribute_type.base_type, 1, 0, attribute_type.role).get_ogn_type_name()}"
        attribute_name = attribute_name.replace("prim", "bundle")
        if attribute_type.tuple_count > 1:
            if attribute_type.role in [og.AttributeRole.TRANSFORM, og.AttributeRole.FRAME, og.AttributeRole.MATRIX]:
                attribute_name += f"_{_MATRIX_DIMENSIONS[attribute_type.tuple_count]}"
            else:
                attribute_name += f"_{attribute_type.tuple_count}"
        array_depth = attribute_type.array_depth
        while array_depth > 0 and attribute_type.role not in [og.AttributeRole.TEXT, og.AttributeRole.PATH]:
            attribute_name += "_array"
            array_depth -= 1
        return attribute_name

    # ----------------------------------------------------------------------
    def __types_to_test(self) -> Set[str]:
        """Returns the set of attribute type names to be tested. This consists of every type accepted by OGN,
        minus a few that are tested separately"""
        # "any" types need a different set of tests as ultimately they only represent one of the other types.
        # "bundle" and "target" types have different tests as they are mainly concerned with membership rather than actual data
        # "transform" types are omitted as they are marked for deprecation in USD so we don't want to support them
        types_not_tested = ["any", "bundle", "target", "transform[4]", "transform[4][]"]
        return [type_name for type_name in ogn.supported_attribute_type_names() if type_name not in types_not_tested]

    # ----------------------------------------------------------------------
    async def __verify_all_types_are_tested(self, node_type_name: str):
        """Meta-test to ensure that the list of types tested in the sample node is the same as the list of all
        supported types, according to the OGN node generator. It will use the full set of supported names,
        drawn from the generator, and compare against the set of attribute type names that appear on the given
        node, extracted from the node's ABI.

        Args:
            node_type_name: Registered node type name for the node to test
        """
        # Get the list of types to test.
        supported_types = self.__types_to_test()

        nodes_attribute_types = []
        controller = og.Controller()
        keys = og.Controller.Keys
        (_, (test_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH, {keys.CREATE_NODES: ("NodeWithAllTypes", node_type_name)}
        )
        for attribute in test_node.get_attributes():
            if attribute.get_port_type() != og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT:
                continue
            attribute_name = attribute.get_name()
            if attribute_name in self.UNCHECKED_ATTRIBUTES:
                continue
            attribute_type = attribute.get_resolved_type()
            port_name = og.get_port_type_namespace(attribute.get_port_type())
            expected_name = f"{port_name}:{self.__attribute_type_to_name(attribute_type)}"
            type_name = attribute_type.get_ogn_type_name()
            self.assertEqual(
                expected_name, attribute_name, f"Generated name differs from actual name, type {type_name}"
            )
            nodes_attribute_types.append(type_name)
        self.assertCountEqual(supported_types, nodes_attribute_types)

    # ----------------------------------------------------------------------
    async def test_all_types_tested_cpp(self):
        """Test that the C++ node exercising all allowed types has covered everything"""
        await self.__verify_all_types_are_tested("omni.graph.test.TestAllDataTypes")

    # ----------------------------------------------------------------------
    async def test_all_types_tested_python(self):
        """Test that the Python node exercising all allowed types has covered everything"""
        await self.__verify_all_types_are_tested("omni.graph.test.TestAllDataTypesPy")

    # ----------------------------------------------------------------------
    async def __test_all_data_types(self, node_type_name: str):
        """Test that sets values of all supported types on the node of the given type using automatic type generators.
        Both the name of the attribute and the sample values for an attribute of that type are supplied automatically.
        The node's only function should be to assign the input to the output, so all types can be validated by simply
        comparing the computed output against the values that were set on the input. It is assumed that actual compute
        operation and input/output correspondence is done elsewhere so this type of test is sufficient to validate that
        access to all data types as both an input and an output is correctly implemented.

        Args:
            node_type_name: Registered node type name for the node to test
        """
        # Get the list of types to test.
        supported_types = self.__types_to_test()

        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (test_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH, {keys.CREATE_NODES: ("NodeWithAllTypes", node_type_name)}
        )
        attribute_names = [
            self.__attribute_type_to_name(og.AttributeType.type_from_ogn_type_name(type_name))
            for type_name in supported_types
        ]
        expected_values = [
            ogn.get_attribute_manager_type(type_name).sample_values()[0] for type_name in supported_types
        ]
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: [
                    ((f"inputs:{name}", test_node), value) for name, value in zip(attribute_names, expected_values)
                ]
                + [(("inputs:doNotCompute", test_node), False)]
            },
        )
        await controller.evaluate(graph)
        actual_values = [
            og.Controller.get(controller.attribute(f"outputs:{name}", test_node)) for name in attribute_names
        ]
        for actual_value, expected_value, attribute_name in zip(actual_values, expected_values, attribute_names):
            ogts.verify_values(expected_value, actual_value, f"Testing output data types on {attribute_name}")
        actual_values = [
            og.Controller.get(controller.attribute(f"state:{name}", test_node)) for name in attribute_names
        ]
        for actual_value, expected_value, attribute_name in zip(actual_values, expected_values, attribute_names):
            ogts.verify_values(expected_value, actual_value, f"Testing state data types on {attribute_name}")

    # ----------------------------------------------------------------------
    async def test_all_data_types_cpp(self):
        """Test that the C++ node exercising all allowed types has covered everything"""
        await self.__test_all_data_types("omni.graph.test.TestAllDataTypes")

    # ----------------------------------------------------------------------
    async def test_all_data_types_python(self):
        """Test that the Python node exercising all allowed types has covered everything"""
        await self.__test_all_data_types("omni.graph.test.TestAllDataTypesPy")

    # ----------------------------------------------------------------------
    async def test_all_types_extended(self):
        """Tests that 'any' attribute can be resolved to and correctly read/write data of all types"""
        carb.log_warn("Warnings of the type 'No source has valid data' are expected during this test")
        # AllTypes resolves the types upstream through A and B. Then we set a value on A, and ensure
        # the value make it through to concrete AllTypes.inputs:attrib_name

        #               'Any'              'Any'
        # +------------+     +------------+     +-----------+
        # | Extended A +---->| Extended B +---->|AllTypes   |
        # +------------+     +------------+     +-----------+

        controller = og.Controller()
        keys = og.Controller.Keys
        (_, (extended_node_a, extended_node_b, all_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("ExtendedNodeA", "omni.graph.test.TypeResolution"),
                    ("ExtendedNodeB", "omni.graph.test.TypeResolution"),
                    ("AllTypes", "omni.graph.test.TestAllDataTypes"),
                ],
                keys.SET_VALUES: ("AllTypes.inputs:doNotCompute", False),
            },
        )

        any_in_a = controller.attribute("inputs:anyValueIn", extended_node_a)
        any_in_b = controller.attribute("inputs:anyValueIn", extended_node_b)
        any_out_a = controller.attribute("outputs:anyValue", extended_node_a)
        any_out_b = controller.attribute("outputs:anyValue", extended_node_b)

        async def test_type(attrib_name, expected_value):
            # Connect the 'any' output to the typed-input to get the 'any' input to be resolved.
            # Then set the value of the resolved input and measure the resolved output value
            (graph, _, _, _) = controller.edit(
                self.TEST_GRAPH_PATH, {keys.CONNECT: [(any_out_b, (attrib_name, all_node)), (any_out_a, any_in_b)]}
            )
            await controller.evaluate(graph)

            controller.edit(self.TEST_GRAPH_PATH, {keys.SET_VALUES: (any_in_a, expected_value)})
            await controller.evaluate(graph)

            actual_value_a = og.Controller.get(any_out_a)
            actual_value_b = og.Controller.get(any_out_b)
            actual_value_all = og.Controller.get(controller.attribute(attrib_name, all_node))

            if isinstance(expected_value, str) or (
                isinstance(expected_value, list) and isinstance(expected_value[0], str)
            ):
                self.assertEqual(expected_value, actual_value_a)
                self.assertEqual(expected_value, actual_value_b)
                self.assertEqual(expected_value, actual_value_all)
            else:
                expected_array = np.ravel(np.array(expected_value))
                self.assertTrue(np.allclose(np.ravel(np.atleast_1d(actual_value_a)), expected_array))
                self.assertTrue(np.allclose(np.ravel(np.atleast_1d(actual_value_b)), expected_array))
                self.assertTrue(np.allclose(np.ravel(np.atleast_1d(actual_value_all)), expected_array))

            # Node totally disconnects - type should be reverted to unknown
            controller.edit(
                self.TEST_GRAPH_PATH, {keys.DISCONNECT: [(any_out_b, (attrib_name, all_node)), (any_out_a, any_in_b)]}
            )
            await controller.evaluate(graph)

        # objectId can not connect to 'any', so remove that from this test
        supported_types = [
            type_name
            for type_name in self.__types_to_test()
            if type_name not in ["objectId", "objectId[]", "execution", "path"]
        ]
        attribute_names = [
            self.__attribute_type_to_name(og.AttributeType.type_from_ogn_type_name(type_name))
            for type_name in supported_types
        ]
        # We can test 2 sets of sample values
        expected_values_set0 = [
            ogn.get_attribute_manager_type(type_name).sample_values()[0] for type_name in supported_types
        ]
        expected_values_set1 = [
            ogn.get_attribute_manager_type(type_name).sample_values()[1] for type_name in supported_types
        ]

        for attrib_name, value in sorted(zip(attribute_names, expected_values_set0)):
            await test_type(f"inputs:{attrib_name}", value)
        for attrib_name, value in sorted(zip(attribute_names, expected_values_set1)):
            await test_type(f"inputs:{attrib_name}", value)

    # ----------------------------------------------------------------------
    async def test_all_usd_cast_types(self):
        """Test that checks that bundles with USD type casting works for all available USD types"""
        # Set up the contents for a prim with all types. Skip types that cannot be added to a prim bundle.
        supported_type_names = [
            type_name
            for type_name in self.__types_to_test()
            if type_name not in ["execution", "objectId", "objectId[]"]
        ]
        supported_types = [og.AttributeType.type_from_ogn_type_name(type_name) for type_name in supported_type_names]
        attribute_names = [self.__attribute_type_to_name(supported_type) for supported_type in supported_types]
        expected_values = [
            ogn.get_attribute_manager_type(type_name).sample_values()[0] for type_name in supported_type_names
        ]
        # This is the array of types for which the node will attempt USD casting, which should be all of the types
        # mentioned in omni/graph/core/ogn/UsdTypes.h
        usd_cast_types = [
            "half",
            "token",
            "timecode",
            "half[]",
            "token[]",
            "timecode[]",
            "double[2]",
            "double[3]",
            "double[4]",
            "double[2][]",
            "double[3][]",
            "double[4][]",
            "float[2]",
            "float[3]",
            "float[4]",
            "float[2][]",
            "float[3][]",
            "float[4][]",
            "half[2]",
            "half[3]",
            "half[4]",
            "half[2][]",
            "half[3][]",
            "half[4][]",
            "int[2]",
            "int[3]",
            "int[4]",
            "int[2][]",
            "int[3][]",
            "int[4][]",
            "matrixd[2]",
            "matrixd[3]",
            "matrixd[4]",
            "matrixd[2][]",
            "matrixd[3][]",
            "matrixd[4][]",
            # These are caught by the timecode cast
            "double",
            "double[]",
            # Roles are caught by the casts too since they are ignored at that point (and not checked in the node)
            "frame[4]",
            "frame[4][]",
            "quatd[4]",
            "quatf[4]",
            "quath[4]",
            "quatd[4][]",
            "quatf[4][]",
            "quath[4][]",
            "colord[3]",
            "colord[3][]",
            "colorf[3]",
            "colorf[3][]",
            "colorh[3]",
            "colorh[3][]",
            "colord[4]",
            "colord[4][]",
            "colorf[4]",
            "colorf[4][]",
            "colorh[4]",
            "colorh[4][]",
            "normald[3]",
            "normald[3][]",
            "normalf[3]",
            "normalf[3][]",
            "normalh[3]",
            "normalh[3][]",
            "pointd[3]",
            "pointd[3][]",
            "pointf[3]",
            "pointf[3][]",
            "pointh[3]",
            "pointh[3][]",
            "vectord[3]",
            "vectord[3][]",
            "vectorf[3]",
            "vectorf[3][]",
            "vectorh[3]",
            "vectorh[3][]",
            "texcoordd[2]",
            "texcoordd[2][]",
            "texcoordf[2]",
            "texcoordf[2][]",
            "texcoordh[2]",
            "texcoordh[2][]",
            "texcoordd[3]",
            "texcoordd[3][]",
            "texcoordf[3]",
            "texcoordf[3][]",
            "texcoordh[3]",
            "texcoordh[3][]",
        ]

        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (_, inspector_outputs_node, inspector_state_node, _), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("UsdCastingNode", "omni.graph.test.TestUsdCasting"),
                    ("InspectorOutput", "omni.graph.nodes.BundleInspector"),
                    ("InspectorState", "omni.graph.nodes.BundleInspector"),
                ],
                keys.CREATE_PRIMS: [
                    (
                        "/World/Prim",
                        {
                            name: (type_name, value)
                            for type_name, name, value in zip(supported_types, attribute_names, expected_values)
                        },
                    )
                ],
                keys.EXPOSE_PRIMS: (og.Controller.PrimExposureType.AS_BUNDLE, "/World/Prim", "PrimExtract"),
                keys.CONNECT: [
                    ("PrimExtract.outputs_primBundle", "UsdCastingNode.inputs:a_bundle"),
                    ("UsdCastingNode.outputs_a_bundle", "InspectorOutput.inputs:bundle"),
                    ("UsdCastingNode.state_a_bundle", "InspectorState.inputs:bundle"),
                ],
            },
        )
        await controller.evaluate(graph)

        expected_output = {}
        count = 1
        for index, type_name in enumerate(supported_type_names):
            if type_name in usd_cast_types:
                count += 1
                expected_output[attribute_names[index]] = [
                    supported_types[index].get_base_type_name(),
                    supported_types[index].tuple_count,
                    supported_types[index].array_depth,
                    supported_types[index].get_role_name(),
                    expected_values[index],
                ]
        # Add in the extra attributes that aren't explicit types
        expected_output["sourcePrimType"] = ["token", 1, 0, "none", "OmniGraphPrim"]
        expected_output["sourcePrimPath"] = ["token", 1, 0, "none", "/World/Prim"]
        expected_results = (len(expected_output), expected_output)

        ognts.verify_bundles_are_equal(
            ognts.filter_bundle_inspector_results(
                ognts.bundle_inspector_results(inspector_outputs_node), [], filter_for_inclusion=False
            ),
            expected_results,
        )

        ognts.verify_bundles_are_equal(
            ognts.filter_bundle_inspector_results(
                ognts.bundle_inspector_results(inspector_state_node), [], filter_for_inclusion=False
            ),
            expected_results,
        )

    # ----------------------------------------------------------------------
    async def __test_read_all_types_default(self, node_type_name: str):
        """Test that saving and loading a file containing non-standard default values correctly brings back the
        default values for all types.

        Args:
            node_type_name: Registered node type name for the node to test
        """
        # Create a file containing a node of the given type and save it
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (test_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH, {keys.CREATE_NODES: ("NodeWithAllTypes", node_type_name)}
        )
        await controller.evaluate(graph)
        test_node_path = test_node.get_prim_path()

        # Get the list of types to test.
        supported_types = self.__types_to_test()

        attribute_names = [
            self.__attribute_type_to_name(og.AttributeType.type_from_ogn_type_name(type_name))
            for type_name in supported_types
        ]

        # Remember the default values in the node. A copy has to be made because otherwise the data will reference
        # Fabric data, which goes away when the stage is saved (because currently at that point OmniGraph is rebuilt).
        expected_default_values = []
        for name in attribute_names:
            value = og.Controller.get(og.Controller.attribute(f"inputs:{name}", test_node))
            if isinstance(value, np.ndarray):
                expected_default_values.append(value.copy())
            else:
                expected_default_values.append(value)

        # --------------------------------------------------------------------------------------------------------------
        # First test that the node that was saved with no values comes back in with the expected default values
        with TemporaryDirectory() as test_directory:
            tmp_file_path = Path(test_directory) / f"{node_type_name}.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)

        await og.Controller.evaluate()
        test_node = og.Controller.node(test_node_path)
        actual_values = [
            og.Controller.get(controller.attribute(f"inputs:{name}", test_node)) for name in attribute_names
        ]
        for name in attribute_names:
            value = og.Controller.get(controller.attribute(f"inputs:{name}", test_node))
        for actual, expected, name in zip(actual_values, expected_default_values, attribute_names):
            ogts.verify_values(expected, actual, f"Testing input data types on {name}")

    # ----------------------------------------------------------------------
    async def test_read_all_types_default_cpp(self):
        """Test that the C++ node exercising all allowed types has covered everything"""
        await self.__test_read_all_types_default("omni.graph.test.TestAllDataTypes")

    # ----------------------------------------------------------------------
    async def test_read_all_types_default_python(self):
        """Test that the Python node exercising all allowed types has covered everything"""
        await self.__test_read_all_types_default("omni.graph.test.TestAllDataTypesPy")

    # ----------------------------------------------------------------------
    async def __test_read_all_types_modified(self, node_type_name: str):
        """Test that saving and loading a file containing non-default values correctly brings back the
        modified values for all types.

        Args:
            node_type_name: Registered node type name for the node to test
        """
        node_type_base_name = node_type_name.split(".")[-1]

        # Get the list of all types to test.
        supported_types = self.__types_to_test()
        attribute_names = [
            self.__attribute_type_to_name(og.AttributeType.type_from_ogn_type_name(type_name))
            for type_name in supported_types
        ]

        # Sample values from the attribute managers are known to not be the defaults of the AllNodeTypes nodes
        expected_sample_values = [
            ogn.get_attribute_manager_type(type_name).sample_values()[0] for type_name in supported_types
        ]

        test_file_name = f"{node_type_base_name}.NonDefault.usda"
        (result, error) = await ogts.load_test_file(test_file_name, use_caller_subdirectory=True)
        self.assertTrue(result, error)

        test_node = og.Controller.node(f"/World/TestGraph/{node_type_base_name}")
        actual_values = [
            og.Controller.get(og.Controller.attribute(f"inputs:{name}", test_node)) for name in attribute_names
        ]
        for actual, expected, name in zip(actual_values, expected_sample_values, attribute_names):
            ogts.verify_values(expected, actual, f"Testing input data types on {name}")

    # ----------------------------------------------------------------------
    async def test_read_all_types_modified_cpp(self):
        """Test that the C++ node exercising all allowed types has covered everything"""
        await self.__test_read_all_types_modified("omni.graph.test.TestAllDataTypes")

    # ----------------------------------------------------------------------
    async def test_read_all_types_modified_python(self):
        """Test that the Python node exercising all allowed types has covered everything"""
        await self.__test_read_all_types_modified("omni.graph.test.TestAllDataTypesPy")

    # ----------------------------------------------------------------------
    def _sanity_check_union_values(self, types: dict):
        """Runs a series of tests to validate the given dictionary of unions contains expected values"""

        def validate_subset(subset: str, superset: str):
            self.assertIsNotNone(types.get(subset, None))
            self.assertIsNotNone(types.get(superset, None))
            subset_items = types[subset]
            superset_items = types[superset]
            for item in subset_items:
                self.assertIn(item, superset_items, f"Expected {item} to be in {superset}")

        def validate_array_form(scalar_form: str, array_form: str):
            self.assertIsNotNone(types.get(scalar_form, None))
            self.assertIsNotNone(types.get(array_form, None))
            scalar_items = types[scalar_form]
            array_items = types[array_form]
            for item in scalar_items:
                self.assertIn(item + "[]", array_items, f"Expected {item} to be in {array_form}")

        self.assertGreater(len(types), 0)

        for x in ["uchar", "int", "uint", "uint64", "int64"]:
            self.assertIn(x, types["integral_scalers"])
        for x in ["int[2]", "int[3]", "int[4]"]:
            self.assertIn(x, types["integral_tuples"])
        for x in ["double", "float", "half", "timecode"]:
            self.assertIn(x, types["decimal_scalers"])
        for x in [
            "double[2]",
            "double[3]",
            "double[4]",
            "float[2]",
            "float[3]",
            "float[4]",
            "half[2]",
            "half[3]",
            "half[4]",
            "colord[3]",
            "colord[4]",
            "colorf[3]",
            "colorf[4]",
            "colorh[3]",
            "colorh[4]",
            "normald[3]",
            "normalf[3]",
            "normalh[3]",
            "pointd[3]",
            "pointf[3]",
            "pointh[3]",
            "texcoordd[2]",
            "texcoordd[3]",
            "texcoordf[2]",
            "texcoordf[3]",
            "texcoordh[2]",
            "texcoordh[3]",
            "quatd[4]",
            "quatf[4]",
            "quath[4]",
            "vectord[3]",
            "vectorf[3]",
            "vectorh[3]",
        ]:
            self.assertIn(x, types["decimal_tuples"])

        validate_subset("integral_scalers", "integral_array_elements")
        validate_subset("integral_tuples", "integral_array_elements")
        validate_array_form("integral_array_elements", "integral_arrays")
        validate_subset("integral_array_elements", "integrals")
        validate_subset("integral_arrays", "integrals")

        validate_subset("decimal_scalers", "decimal_array_elements")
        validate_subset("decimal_tuples", "decimal_array_elements")
        validate_array_form("decimal_array_elements", "decimal_arrays")
        validate_subset("decimal_array_elements", "decimals")
        validate_subset("decimal_arrays", "decimals")

        validate_subset("integral_scalers", "numeric_scalers")
        validate_subset("decimal_scalers", "numeric_scalers")
        validate_subset("integral_tuples", "numeric_tuples")
        validate_subset("decimal_tuples", "numeric_tuples")
        validate_subset("numeric_scalers", "numeric_array_elements")
        validate_subset("numeric_tuples", "numeric_array_elements")
        validate_subset("matrices", "numeric_array_elements")
        validate_array_form("numeric_array_elements", "numeric_arrays")

        validate_subset("numeric_array_elements", "numerics")
        validate_subset("numeric_arrays", "numerics")
        validate_subset("numeric_array_elements", "array_elements")
        validate_subset("numeric_arrays", "arrays")

        self.assertIn("token", types["array_elements"])
        self.assertIn("token[]", types["arrays"])

    # ----------------------------------------------------------------------
    async def test_union_types_contain_expected_values(self):
        """Performs a sanity check on the attribute union values returned from the runtime"""

        # The same configuration file is used to populate dictionaries form both dlls.
        # Validate they contain the same entries

        # Test the dictionary loaded from omni.graph.core
        self._sanity_check_union_values(og.AttributeType.get_unions())

        # Validate the entries are the same
        self.assertDictEqual(og.AttributeType.get_unions(), ogn.ATTRIBUTE_UNION_GROUPS)

    # ----------------------------------------------------------------------
    async def __test_nan_values(self, node_type_name: str):
        """Test that evaluating NaN values correctly copies and extracts the values.
        This could not be done with a regular test sequence as NaN != NaN

        Args:
            node_type_name: Registered node type name for the node to test
        """
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (test_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH, {keys.CREATE_NODES: ("NodeWithNans", node_type_name)}
        )
        await controller.evaluate(graph)
        nan_attributes = [
            "colord3",
            "colord4",
            "colord3_array",
            "colord4_array",
            "colorf3",
            "colorf4",
            "colorf3_array",
            "colorf4_array",
            "colorh3",
            "colorh4",
            "colorh3_array",
            "colorh4_array",
            "double",
            "double2",
            "double3",
            "double4",
            "double_array",
            "double2_array",
            "double3_array",
            "double4_array",
            "float",
            "float2",
            "float3",
            "float4",
            "float_array",
            "float2_array",
            "float3_array",
            "float4_array",
            "frame4",
            "frame4_array",
            "half",
            "half2",
            "half3",
            "half4",
            "half_array",
            "half2_array",
            "half3_array",
            "half4_array",
            "matrixd2",
            "matrixd3",
            "matrixd4",
            "matrixd2_array",
            "matrixd3_array",
            "matrixd4_array",
            "normald3",
            "normald3_array",
            "normalf3",
            "normalf3_array",
            "normalh3",
            "normalh3_array",
            "pointd3",
            "pointd3_array",
            "pointf3",
            "pointf3_array",
            "pointh3",
            "pointh3_array",
            "quatd4",
            "quatd4_array",
            "quatf4",
            "quatf4_array",
            "quath4",
            "quath4_array",
            "texcoordd2",
            "texcoordd3",
            "texcoordd2_array",
            "texcoordd3_array",
            "texcoordf2",
            "texcoordf3",
            "texcoordf2_array",
            "texcoordf3_array",
            "texcoordh2",
            "texcoordh3",
            "texcoordh2_array",
            "texcoordh3_array",
            "timecode",
            "timecode_array",
            "vectord3",
            "vectord3_array",
            "vectorf3",
            "vectorf3_array",
            "vectorh3",
            "vectorh3_array",
        ]

        def __is_nan(value) -> bool:
            if isinstance(value, list):
                return all(__is_nan(element) for element in value)
            if isinstance(value, np.ndarray):
                return np.alltrue(np.isnan(value))
            return isnan(value)

        for attribute_name in nan_attributes:
            attr = controller.attribute(f"outputs:a_{attribute_name}_nan", test_node)
            value = controller.get(attr)
            self.assertTrue(__is_nan(value))

            attr = controller.attribute(f"outputs:a_{attribute_name}_snan", test_node)
            value = controller.get(attr)
            self.assertTrue(__is_nan(value))

    # ----------------------------------------------------------------------
    async def test_nan_values_cpp(self):
        """Test that the C++ node exercising all allowed types handles NaN values correctly"""
        await self.__test_nan_values("omni.graph.test.TestNanInf")

    # ----------------------------------------------------------------------
    async def test_nan_values_python(self):
        """Test that the Python node exercising all allowed types handles NaN values correctly"""
        await self.__test_nan_values("omni.graph.test.TestNanInfPy")

    async def test_node_database_id_py(self):
        """Test if database id is not changing between graph evaluations"""
        controller = og.Controller()
        keys = og.Controller.Keys

        # Disconnected input bundle should be invalid
        (graph, (_, node_database), _, _) = controller.edit(
            "/World/Graph",
            {
                keys.CREATE_NODES: [
                    ("BundleConstructor", "omni.graph.nodes.BundleConstructor"),
                    ("NodeDatabase", "omni.graph.test.NodeDatabasePy"),
                ],
            },
        )

        # database id remains the same because topology has not changed
        graph.evaluate()
        previous_id = node_database.get_attribute("outputs:id").get()
        graph.evaluate()
        current_id = node_database.get_attribute("outputs:id").get()
        self.assertTrue(previous_id == current_id)

        previous_id = current_id
        graph.evaluate()
        current_id = node_database.get_attribute("outputs:id").get()
        self.assertTrue(previous_id == current_id)

        # disconnect input connect rebuilds the database
        (graph, _, _, _) = controller.edit(
            "/World/Graph",
            {
                keys.CONNECT: [
                    ("BundleConstructor.outputs_bundle", "NodeDatabase.inputs:bundle"),
                ],
            },
        )

        graph.evaluate()
        current_id = node_database.get_attribute("outputs:id").get()
        self.assertTrue(previous_id != current_id)

        # subsequent evaluation does not change id
        previous_id = current_id
        graph.evaluate()
        current_id = node_database.get_attribute("outputs:id").get()
        self.assertTrue(previous_id == current_id)
