"""
Tests for attribute bundles
"""
import unittest
from contextlib import suppress
from typing import Any, List, Tuple

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.graph.nodes.tests as ognts
import omni.graph.tools as ogt
from omni.kit.test.teamcity import is_running_in_teamcity
from pxr import Gf


# ==============================================================================================================
def multiply_elements(multiplier, original: Any) -> Any:
    """Take in a numeric value, tuple, list, tuple of tuples, etc. and multiply every element by a constant"""
    if isinstance(original, tuple):
        return tuple(multiply_elements(multiplier, element) for element in original)
    if isinstance(original, list):
        return tuple(multiply_elements(multiplier, element) for element in original)
    return multiplier * original


# ==============================================================================================================
def as_tuple(value: Any) -> Tuple:
    """Return the value, interpreted as a tuple (except simple values, which are returned as themselves)"""
    if isinstance(value, tuple):
        return value
    if isinstance(value, list):
        return tuple(value)
    if isinstance(value, (Gf.Quatd, Gf.Quatf, Gf.Quath)):
        return (*value.GetImaginary(), value.GetReal())
    if isinstance(value, Gf.Matrix4d):
        return tuple(tuple(element for element in list(row)) for row in value)
    with suppress(TypeError):
        if len(value) > 1:
            return tuple(value)
    return value


# ==============================================================================================================
class TestOmniGraphBundles(ogts.OmniGraphTestCase):
    """Attribute bundles do not yet have the ability to log tests directly in a .ogn file so it's done here"""

    # --------------------------------------------------------------------------------------------------------------
    def _compare_results(self, expected_raw: Any, actual_raw: Any, test_info: str):
        """Loose comparison of two types of compatible values

        Args:
            expected_raw: Expected results. Can be simple values, tuples, lists, or pxr.Gf types
            actual_raw: Actual results. Can be simple values, tuples, lists, or pxr.Gf types
            test_info: String to accompany error messages

        Returns:
            Error encountered when mismatched values found, None if everything matched
        """
        expected = as_tuple(expected_raw)
        actual = as_tuple(actual_raw)
        self.assertEqual(
            type(expected), type(actual), f"{test_info} Mismatched types - expected {expected_raw}, got {actual_raw}"
        )

        if isinstance(expected, tuple):
            for expected_element, actual_element in zip(expected, actual):
                self._compare_results(expected_element, actual_element, test_info)
        else:
            self.assertEqual(expected, actual, f"{test_info} Expected {expected}, got {actual}")

    # ----------------------------------------------------------------------
    async def _test_bundle_contents(self, node_type_to_test: str):
        """Test access to bundle attribute manipulation for C++ and Python implementations"""
        keys = og.Controller.Keys
        (graph, [bundle_node, inspector_node, _, _], [_, filtered_prim], _) = og.Controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("BundleManipulator", node_type_to_test),
                    ("Inspector", "omni.graph.nodes.BundleInspector"),
                ],
                keys.CREATE_PRIMS: [
                    ("FullPrim", {"fullInt": ("int", 2), "floatArray": ("float[]", [3.0, 4.0, 5.0])}),
                    (
                        "FilteredPrim",
                        {
                            "int_1": ("int", 6),
                            "uint_1": ("uint", 7),
                            "int64_x": ("int64", 8),
                            "uint64_1": ("uint64", 9),
                            "int_3": ("int[3]", (10, 11, 12)),
                            "int_array": ("int[]", [13, 14, 15]),
                            "float_1": ("float", 3.0),
                            "double_x": ("double", 4.0),
                            "big_int": ("int[]", [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]),
                        },
                    ),
                ],
                keys.EXPOSE_PRIMS: [
                    (og.Controller.PrimExposureType.AS_BUNDLE, "/FullPrim", "FullPrimExtract"),
                    (og.Controller.PrimExposureType.AS_BUNDLE, "/FilteredPrim", "FilteredPrimExtract"),
                ],
                keys.CONNECT: [
                    ("FullPrimExtract.outputs_primBundle", "BundleManipulator.inputs:fullBundle"),
                    ("FilteredPrimExtract.outputs_primBundle", "BundleManipulator.inputs:filteredBundle"),
                    ("BundleManipulator.outputs_combinedBundle", "Inspector.inputs:bundle"),
                ],
            },
        )
        _ = ogt.OGN_DEBUG and ognts.enable_debugging(inspector_node)
        await og.Controller.evaluate()

        expected_results = {
            "big_int": ("int", 1, 1, "none", [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]),
            "int_1": ("int", 1, 0, "none", 6),
            "uint_1": ("uint", 1, 0, "none", 7),
            "int64_x": ("int64", 1, 0, "none", 8),
            "uint64_1": ("uint64", 1, 0, "none", 9),
            "int_3": ("int", 3, 0, "none", (10, 11, 12)),
            "int_array": ("int", 1, 1, "none", [13, 14, 15]),
            "float_1": ("float", 1, 0, "none", 3.0),
            "double_x": ("double", 1, 0, "none", 4.0),
            "fullInt": ("int", 1, 0, "none", 2),
            "floatArray": ("float", 1, 1, "none", [3.0, 4.0, 5.0]),
            "sourcePrimPath": ("token", 1, 0, "none", str(filtered_prim.GetPrimPath())),
            "sourcePrimType": ("token", 1, 0, "none", str(filtered_prim.GetTypeName())),
        }

        def __result_subset(elements_removed: List[str]):
            """Return the set of results with the index list of elements removed"""
            return {key: value for key, value in expected_results.items() if key not in elements_removed}

        # Test data consists of a list of individual test configurations with the filters to set and the
        # expected contents of the output bundle using those filters.
        test_data = [
            ([], expected_results),
            (["x"], __result_subset(["int64_x", "double_x"])),
            (["int"], __result_subset(["big_int", "int_1", "uint_1", "int64_x", "uint64_1", "int_3", "int_array"])),
            (["big"], __result_subset(["big_int"])),
            (
                ["x", "big", "int"],
                __result_subset(
                    ["big_int", "int_1", "uint_1", "int64_x", "uint64_1", "int_3", "int_array", "double_x"]
                ),
            ),
        ]

        for (filters, expected_results) in test_data:
            og.Controller.edit(
                graph,
                {
                    keys.SET_VALUES: (("inputs:filters", bundle_node), filters),
                },
            )
            await og.Controller.evaluate()
            try:
                ognts.verify_bundles_are_equal(
                    ognts.filter_bundle_inspector_results(
                        ognts.bundle_inspector_results(inspector_node), [], filter_for_inclusion=False
                    ),
                    ognts.filter_bundle_inspector_results(
                        (len(expected_results), expected_results), [], filter_for_inclusion=False
                    ),
                )
            except ValueError as error:
                self.assertTrue(False, error)

    # ----------------------------------------------------------------------
    async def test_bundle_contents_cpp(self):
        """Test access to bundle attribute manipulation on the C++ implemented node"""
        await self._test_bundle_contents("omni.graph.tutorials.BundleManipulation")

    # ----------------------------------------------------------------------
    async def test_bundle_contents_py(self):
        """Test bundle attribute manipulation on the Python implemented node"""
        await self._test_bundle_contents("omni.graph.tutorials.BundleManipulationPy")

    # ----------------------------------------------------------------------
    async def _test_simple_bundled_data(self, node_type_to_test: str):
        """Test access to attributes with simple data types within bundles for both C++ and Python implementations"""
        controller = og.Controller()
        keys = og.Controller.Keys
        prim_definition = ognts.prim_with_everything_definition()

        (_, [_, inspector_node, _], _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("BundledDataModifier", node_type_to_test),
                    ("Inspector", "omni.graph.nodes.BundleInspector"),
                ],
                keys.CREATE_PRIMS: ("TestPrim", prim_definition),
                keys.EXPOSE_PRIMS: (og.Controller.PrimExposureType.AS_BUNDLE, "TestPrim", "TestBundle"),
                keys.CONNECT: [
                    ("TestBundle.outputs_primBundle", "BundledDataModifier.inputs:bundle"),
                    ("BundledDataModifier.outputs_bundle", "Inspector.inputs:bundle"),
                ],
            },
        )

        _ = ogt.OGN_DEBUG and ognts.enable_debugging(inspector_node)
        await controller.evaluate()

        (bundled_count, bundled_results) = ognts.bundle_inspector_results(inspector_node)
        # Prim will not have the "sourcePrimXX" attributes so it is smaller by 2
        self.assertEqual(len(prim_definition) + 2, bundled_count)

        for name, (ogn_type, array_depth, tuple_count, role, actual_output) in bundled_results.items():
            if name in ["sourcePrimType", "sourcePrimPath"]:
                continue
            if not node_type_to_test.endswith("Py") and ogn_type not in ["int", "int64", "uint", "uint64"]:
                # The C++ node only handles integer types, in the interest of brevity
                continue
            test_info = "Bundled"
            if tuple_count > 1:
                test_info += f" tuple[{tuple_count}]"
            if array_depth > 0:
                test_info += " array"
            test_info += f" attribute {name} of type {ogn_type}"
            if role != "none":
                test_info += f" (role {role})"
            if ogn_type == "token" or role == "text":
                expected_output = prim_definition[name][1]
                if isinstance(expected_output, str):
                    expected_output += expected_output
                else:
                    expected_output = [f"{element}{element}" for element in expected_output]
                self._compare_results(expected_output, actual_output, test_info)
            else:
                if ogn_type in ["bool", "bool[]"]:
                    expected_output = as_tuple(prim_definition[name][1])
                else:
                    expected_output = multiply_elements(2, as_tuple(prim_definition[name][1]))
                self._compare_results(expected_output, actual_output, test_info)

    # ----------------------------------------------------------------------
    async def test_simple_bundled_data_cpp(self):
        """Test access to attributes with simple data types within bundles on the C++ implemented node"""
        await self._test_simple_bundled_data("omni.graph.tutorials.BundleData")

    # ----------------------------------------------------------------------
    async def test_simple_bundled_data_py(self):
        """Test access to attributes with simple data types within bundles on the Python implemented node"""
        await self._test_simple_bundled_data("omni.graph.tutorials.BundleDataPy")

    # ----------------------------------------------------------------------
    async def _test_add_attributes_to_bundle(self, node_type_to_test: str):
        """Test basic operation of the tutorial node that adds new attributes to a bundle"""

        keys = og.Controller.Keys
        (graph, [add_node, inspector_node], _, _) = og.Controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("AddAttributes", node_type_to_test),
                    ("Inspector", "omni.graph.nodes.BundleInspector"),
                ],
                keys.CONNECT: ("AddAttributes.outputs_bundle", "Inspector.inputs:bundle"),
            },
        )
        await og.Controller.evaluate()

        types_attribute = og.Controller.attribute("inputs:typesToAdd", add_node)
        added_names_attribute = og.Controller.attribute("inputs:addedAttributeNames", add_node)
        removed_names_attribute = og.Controller.attribute("inputs:removedAttributeNames", add_node)

        _ = ogt.OGN_DEBUG and ognts.enable_debugging(inspector_node)

        # List of test data configurations to run. The individual test entries consist of:
        #    - list of the types for the attributes to be added
        #    - list of the base types corresponding to the main types
        #    - list of names for the attributes to be added
        #    - list of values expected for the new bundle as (count, roles, arrayDepths, tupleCounts, values)
        test_data = [
            [
                ["float", "double", "int", "int64", "uchar", "uint", "uint64", "bool", "token"],
                ["float", "double", "int", "int64", "uchar", "uint", "uint64", "bool", "token"],
                [f"output{n}" for n in range(9)],
                (9, ["none"] * 9, [0] * 9, [1] * 9, [0.0, 0.0, 0, 0, 0, 0, 0, False, ""]),
            ],
            [
                ["float[]", "double[]", "int[]", "int64[]", "uchar[]", "uint[]", "uint64[]", "token[]"],
                ["float", "double", "int", "int64", "uchar", "uint", "uint64", "token"],
                [f"output{n}" for n in range(8)],
                (8, ["none"] * 8, [1] * 8, [1] * 8, [[], [], [], [], [], [], [], []]),
            ],
            [
                ["float[3]", "double[2]", "int[4]"],
                ["float", "double", "int"],
                [f"output{n}" for n in range(3)],
                (3, ["none"] * 3, [0] * 3, [3, 2, 4], [[0.0, 0.0, 0.0], [0.0, 0.0], [0, 0, 0, 0]]),
            ],
            [
                ["float[3][]", "double[2][]", "int[4][]"],
                ["float", "double", "int"],
                [f"output{n}" for n in range(3)],
                (3, ["none"] * 3, [1] * 3, [3, 2, 4], [[], [], []]),
            ],
            [
                ["any"],
                ["token"],
                ["output0"],
                (1, ["none"], [0], [1], [""]),
            ],
            [
                ["colord[3]", "colorf[4]", "colorh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["color"] * 3, [0] * 3, [3, 4, 3], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["normald[3]", "normalf[3]", "normalh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["normal"] * 3, [0] * 3, [3, 3, 3], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["pointd[3]", "pointf[3]", "pointh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["point"] * 3, [0] * 3, [3, 3, 3], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["quatd[4]", "quatf[4]", "quath[4]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (
                    3,
                    ["quat"] * 3,
                    [0] * 3,
                    [4, 4, 4],
                    [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]],
                ),
            ],
            [
                ["texcoordd[3]", "texcoordf[2]", "texcoordh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["texcoord"] * 3, [0] * 3, [3, 2, 3], [[0.0, 0.0, 0.0], [0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["vectord[3]", "vectorf[3]", "vectorh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["vector"] * 3, [0] * 3, [3, 3, 3], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["timecode"],
                ["double"],
                ["output0"],
                (1, ["timecode"], [0], [1], [0.0]),
            ],
        ]

        for (types, base_types, names, results) in test_data:
            og.Controller.edit(
                graph,
                {
                    keys.SET_VALUES: [
                        (types_attribute, types),
                        (added_names_attribute, names),
                    ]
                },
            )
            await og.Controller.evaluate()

            (bundled_count, bundled_results) = ognts.bundle_inspector_results(inspector_node)
            self.assertEqual(bundled_count, results[0])
            self.assertCountEqual(list(bundled_results.keys()), names)
            for index, name in enumerate(names):
                error = f"Checking {{}} for attribute {name}"
                named_results = bundled_results[name]
                self.assertEqual(
                    named_results[ognts.BundleResultKeys.TYPE_IDX], base_types[index], error.format("type")
                )
                self.assertEqual(
                    named_results[ognts.BundleResultKeys.ARRAY_DEPTH_IDX],
                    results[2][index],
                    error.format("array depth"),
                )
                self.assertEqual(
                    named_results[ognts.BundleResultKeys.TUPLE_COUNT_IDX],
                    results[3][index],
                    error.format("tuple count"),
                )
                self.assertEqual(
                    named_results[ognts.BundleResultKeys.ROLE_IDX], results[1][index], error.format("role")
                )

        # One final test run that also includes a remove of the 2nd, 4th, and 7th attributes in the first test list
        (types, base_types, names, results) = test_data[0]
        og.Controller.edit(
            graph,
            {
                keys.SET_VALUES: [
                    (types_attribute, types),
                    (added_names_attribute, names),
                    (removed_names_attribute, [names[2], names[4], names[7]]),
                ]
            },
        )
        await og.Controller.evaluate()

        def pruned_list(original: List) -> List:
            """Remove the elements 2, 4, 7 from the list and return the result"""
            return [item for index, item in enumerate(original) if index not in [2, 4, 7]]

        (bundled_count, bundled_results) = ognts.bundle_inspector_results(inspector_node)

        # Modify the expected results to account for the removed attributes
        self.assertEqual(bundled_count, results[0] - 3)
        names_expected = pruned_list(names)
        base_types_expected = pruned_list(base_types)
        array_depths_expected = pruned_list(results[2])
        tuple_counts_expected = pruned_list(results[3])
        roles_expected = pruned_list(results[1])

        self.assertCountEqual(list(bundled_results.keys()), names_expected)
        for index, name in enumerate(names_expected):
            error = f"Checking {{}} for attribute {name}"
            named_results = bundled_results[name]
            self.assertEqual(
                named_results[ognts.BundleResultKeys.TYPE_IDX], base_types_expected[index], error.format("type")
            )
            self.assertEqual(
                named_results[ognts.BundleResultKeys.ARRAY_DEPTH_IDX],
                array_depths_expected[index],
                error.format("array depth"),
            )
            self.assertEqual(
                named_results[ognts.BundleResultKeys.TUPLE_COUNT_IDX],
                tuple_counts_expected[index],
                error.format("tuple count"),
            )
            self.assertEqual(
                named_results[ognts.BundleResultKeys.ROLE_IDX], roles_expected[index], error.format("role")
            )

    # ----------------------------------------------------------------------
    async def test_add_attributes_to_bundle_cpp(self):
        """Test adding attributes to bundles on the C++ implemented node"""
        await self._test_add_attributes_to_bundle("omni.graph.tutorials.BundleAddAttributes")

    # ----------------------------------------------------------------------
    async def test_add_attributes_to_bundle_py(self):
        """Test adding attributes to bundles on the Python implemented node"""
        await self._test_add_attributes_to_bundle("omni.graph.tutorials.BundleAddAttributesPy")

    # ----------------------------------------------------------------------
    async def _test_batched_add_attributes_to_bundle(self, node_type_to_test: str):
        """Test basic operation of the tutorial node that adds new attributes to a bundle"""

        keys = og.Controller.Keys
        (graph, [add_node, inspector_node], _, _) = og.Controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [
                    ("AddAttributes", node_type_to_test),
                    ("Inspector", "omni.graph.nodes.BundleInspector"),
                ],
                keys.CONNECT: ("AddAttributes.outputs_bundle", "Inspector.inputs:bundle"),
            },
        )
        await og.Controller.evaluate()

        types_attribute = og.Controller.attribute("inputs:typesToAdd", add_node)
        added_names_attribute = og.Controller.attribute("inputs:addedAttributeNames", add_node)
        removed_names_attribute = og.Controller.attribute("inputs:removedAttributeNames", add_node)
        batched_api_attribute = og.Controller.attribute("inputs:useBatchedAPI", add_node)

        _ = ogt.OGN_DEBUG and ognts.enable_debugging(inspector_node)

        # List of test data configurations to run. The individual test entries consist of:
        #    - list of the types for the attributes to be added
        #    - list of the base types corresponding to the main types
        #    - list of names for the attributes to be added
        #    - list of values expected for the new bundle as (count, roles, arrayDepths, tupleCounts, values)
        test_data = [
            [
                ["float", "double", "int", "int64", "uchar", "uint", "uint64", "bool", "token"],
                ["float", "double", "int", "int64", "uchar", "uint", "uint64", "bool", "token"],
                [f"output{n}" for n in range(9)],
                (9, ["none"] * 9, [0] * 9, [1] * 9, [0.0, 0.0, 0, 0, 0, 0, 0, False, ""]),
            ],
            [
                ["float[]", "double[]", "int[]", "int64[]", "uchar[]", "uint[]", "uint64[]", "token[]"],
                ["float", "double", "int", "int64", "uchar", "uint", "uint64", "token"],
                [f"output{n}" for n in range(8)],
                (8, ["none"] * 8, [1] * 8, [1] * 8, [[], [], [], [], [], [], [], []]),
            ],
            [
                ["float[3]", "double[2]", "int[4]"],
                ["float", "double", "int"],
                [f"output{n}" for n in range(3)],
                (3, ["none"] * 3, [0] * 3, [3, 2, 4], [[0.0, 0.0, 0.0], [0.0, 0.0], [0, 0, 0, 0]]),
            ],
            [
                ["float[3][]", "double[2][]", "int[4][]"],
                ["float", "double", "int"],
                [f"output{n}" for n in range(3)],
                (3, ["none"] * 3, [1] * 3, [3, 2, 4], [[], [], []]),
            ],
            [
                ["any"],
                ["token"],
                ["output0"],
                (1, ["none"], [0], [1], [""]),
            ],
            [
                ["colord[3]", "colorf[4]", "colorh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["color"] * 3, [0] * 3, [3, 4, 3], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["normald[3]", "normalf[3]", "normalh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["normal"] * 3, [0] * 3, [3, 3, 3], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["pointd[3]", "pointf[3]", "pointh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["point"] * 3, [0] * 3, [3, 3, 3], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["quatd[4]", "quatf[4]", "quath[4]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (
                    3,
                    ["quat"] * 3,
                    [0] * 3,
                    [4, 4, 4],
                    [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]],
                ),
            ],
            [
                ["texcoordd[3]", "texcoordf[2]", "texcoordh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["texcoord"] * 3, [0] * 3, [3, 2, 3], [[0.0, 0.0, 0.0], [0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["vectord[3]", "vectorf[3]", "vectorh[3]"],
                ["double", "float", "half"],
                [f"output{n}" for n in range(3)],
                (3, ["vector"] * 3, [0] * 3, [3, 3, 3], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
            ],
            [
                ["timecode"],
                ["double"],
                ["output0"],
                (1, ["timecode"], [0], [1], [0.0]),
            ],
        ]

        for (types, base_types, names, results) in test_data:
            og.Controller.edit(
                graph,
                {
                    keys.SET_VALUES: [
                        (types_attribute, types),
                        (added_names_attribute, names),
                        (batched_api_attribute, True),
                    ]
                },
            )
            await og.Controller.evaluate()

            (bundled_count, bundled_results) = ognts.bundle_inspector_results(inspector_node)
            self.assertEqual(bundled_count, results[0])
            self.assertCountEqual(list(bundled_results.keys()), names)
            for index, name in enumerate(names):
                error = f"Checking {{}} for attribute {name}"
                named_results = bundled_results[name]
                self.assertEqual(
                    named_results[ognts.BundleResultKeys.TYPE_IDX], base_types[index], error.format("type")
                )
                self.assertEqual(
                    named_results[ognts.BundleResultKeys.ARRAY_DEPTH_IDX],
                    results[2][index],
                    error.format("array depth"),
                )
                self.assertEqual(
                    named_results[ognts.BundleResultKeys.TUPLE_COUNT_IDX],
                    results[3][index],
                    error.format("tuple count"),
                )
                self.assertEqual(
                    named_results[ognts.BundleResultKeys.ROLE_IDX], results[1][index], error.format("role")
                )

        # One final test run that also includes a remove of the 2nd, 4th, and 7th attributes in the first test list
        (types, base_types, names, results) = test_data[0]
        og.Controller.edit(
            graph,
            {
                keys.SET_VALUES: [
                    (types_attribute, types),
                    (added_names_attribute, names),
                    (removed_names_attribute, [names[2], names[4], names[7]]),
                ]
            },
        )
        await og.Controller.evaluate()

        def pruned_list(original: List) -> List:
            """Remove the elements 2, 4, 7 from the list and return the result"""
            return [item for index, item in enumerate(original) if index not in [2, 4, 7]]

        await og.Controller.evaluate()

        (bundled_count, bundled_results) = ognts.bundle_inspector_results(inspector_node)
        self.assertEqual(bundled_count, results[0] - 3)
        # Modify the expected results to account for the removed attributes
        names_expected = pruned_list(names)
        base_types_expected = pruned_list(base_types)
        array_depths_expected = pruned_list(results[2])
        tuple_counts_expected = pruned_list(results[3])
        roles_expected = pruned_list(results[1])

        self.assertCountEqual(list(bundled_results.keys()), names_expected)
        for index, name in enumerate(names_expected):
            error = f"Checking {{}} for attribute {name}"
            named_results = bundled_results[name]
            self.assertEqual(
                named_results[ognts.BundleResultKeys.TYPE_IDX], base_types_expected[index], error.format("type")
            )
            self.assertEqual(
                named_results[ognts.BundleResultKeys.ARRAY_DEPTH_IDX],
                array_depths_expected[index],
                error.format("array depth"),
            )
            self.assertEqual(
                named_results[ognts.BundleResultKeys.TUPLE_COUNT_IDX],
                tuple_counts_expected[index],
                error.format("tuple count"),
            )
            self.assertEqual(
                named_results[ognts.BundleResultKeys.ROLE_IDX], roles_expected[index], error.format("role")
            )

    # ----------------------------------------------------------------------
    async def test_batched_add_attributes_to_bundle_cpp(self):
        """Test adding attributes to bundles on the C++ implemented node"""
        await self._test_batched_add_attributes_to_bundle("omni.graph.tutorials.BundleAddAttributes")

    # ----------------------------------------------------------------------
    async def test_batched_add_attributes_to_bundle_py(self):
        """Test adding attributes to bundles on the Python implemented node"""
        await self._test_batched_add_attributes_to_bundle("omni.graph.tutorials.BundleAddAttributesPy")

    # ----------------------------------------------------------------------
    async def _test_bundle_gpu(self, node_type_to_test: str):
        """Test basic operation of the tutorial node that accesses bundle data on the GPU"""

        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (bundle_node, inspector_node), _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_PRIMS: [
                    ("Prim1", {"points": ("pointf[3][]", [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])}),
                    ("Prim2", {"points": ("pointf[3][]", [[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])}),
                ],
                keys.CREATE_NODES: [
                    ("GpuBundleNode", node_type_to_test),
                    ("Inspector", "omni.graph.nodes.BundleInspector"),
                ],
            },
        )

        controller.edit(
            graph,
            {
                keys.EXPOSE_PRIMS: [
                    (og.GraphController.PrimExposureType.AS_BUNDLE, "Prim1", "Prim1Extract"),
                    (og.GraphController.PrimExposureType.AS_BUNDLE, "Prim2", "Prim2Extract"),
                ],
                keys.CONNECT: [
                    ("Prim1Extract.outputs_primBundle", "GpuBundleNode.inputs:cpuBundle"),
                    ("Prim2Extract.outputs_primBundle", "GpuBundleNode.inputs:gpuBundle"),
                    ("GpuBundleNode.outputs_cpuGpuBundle", "Inspector.inputs:bundle"),
                ],
            },
        )

        _ = ogt.OGN_DEBUG and ognts.enable_debugging(inspector_node)
        await controller.evaluate(graph)

        for gpu in [False, True]:
            controller.attribute("inputs:gpu", bundle_node, graph).set(gpu)
            await controller.evaluate()
            (_, bundled_values) = ognts.bundle_inspector_results(inspector_node)
            self.assertEqual(bundled_values["dotProducts"][ognts.BundleResultKeys.VALUE_IDX], [32.0, 122.0])

    # ----------------------------------------------------------------------
    async def test_bundle_gpu_cpp(self):
        """Test access to bundle contents on GPU on the C++ implemented node"""
        await self._test_bundle_gpu("omni.graph.tutorials.CpuGpuBundles")

    # ----------------------------------------------------------------------
    async def test_bundle_gpu_py(self):
        """Test bundle contents on GPU on the Python implemented node"""
        await self._test_bundle_gpu("omni.graph.tutorials.CpuGpuBundlesPy")

    # ----------------------------------------------------------------------
    @unittest.skipIf(is_running_in_teamcity(), "This is a manual test so it only needs to run locally")
    async def test_cuda_pointers_py(self):
        """Run a simple test on a node that extracts CUDA pointers from the GPU. Inspect the output for information"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (_, inspector_node), _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_PRIMS: [
                    ("Prim1", {"points": ("float[3][]", [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])}),
                ],
                keys.CREATE_NODES: [
                    ("GpuBundleNode", "omni.graph.tutorials.CudaCpuArraysPy"),
                    ("Inspector", "omni.graph.nodes.BundleInspector"),
                ],
            },
        )

        controller.edit(
            graph,
            {
                keys.EXPOSE_PRIMS: [
                    (og.GraphController.PrimExposureType.AS_ATTRIBUTES, "Prim1", "Prim1Extract"),
                ],
            },
        )

        _ = ogt.OGN_DEBUG and ognts.enable_debugging(inspector_node)
        await controller.evaluate(graph)

        controller.edit(
            graph,
            {
                keys.CONNECT: [
                    ("Prim1Extract.outputs:points", "GpuBundleNode.inputs:points"),
                    ("GpuBundleNode.outputs_outBundle", "Inspector.inputs:bundle"),
                ],
            },
        )
