# noqa: PLC0302
"""Unit tests for the OgnWritePrims node in this extension"""

import fnmatch
import os
import tempfile

import carb
import omni.graph.core as og
import omni.graph.core.autonode as oga
import omni.graph.core.tests as ogts
import omni.kit.commands
import omni.kit.stage_templates
import omni.kit.test
import omni.timeline
import omni.usd
from pxr import Sdf, Usd, UsdGeom
from usdrt import Sdf as SdfRT
from usdrt import Usd as UsdRT

# Define a test node that has one relationship in the bundle
# This will be used to test WritePrimsV2 node's Relationship exporting
REL_BUNDLE_TRL_NAME = "test_rel"
REL_BUNDLE_TARGETS = [Sdf.Path("/Foo/bar"), Sdf.Path("/Foo/baz")]
REL_BUNDLE_TARGETS_RT = [SdfRT.Path(path.pathString) for path in REL_BUNDLE_TARGETS]


@oga.AutoFunc(module_name="__test_write_prims_v2__", pure=True)
def relationship_bundle() -> og.Bundle:
    bundle = og.Bundle("return", False)
    bundle.create_attribute(REL_BUNDLE_TRL_NAME, og.Type(og.BaseDataType.RELATIONSHIP, 1, 1)).value = REL_BUNDLE_TARGETS
    return bundle


# ======================================================================


# because the child bundle's order inside of bundle is non-deterministic, WritePrims node sort the bundle content
# according to sourcePrimPath when scatter_under_targets mode is off
def sort_child_bundles(child_bundles, scatter_under_targets: bool):
    # do not sort if scatter_under_targets
    if scatter_under_targets:
        return child_bundles

    child_bundles.sort(key=lambda x: x.get_attribute_by_name("sourcePrimPath").get())
    return child_bundles


# ======================================================================


class TestWritePrimNodes(ogts.OmniGraphTestCase):
    TEST_GRAPH_PATH = "/World/TestGraph"

    # ----------------------------------------------------------------------

    async def test_read_write_prims_mpib(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB without bundle modification and output to targets"""
        await self._test_read_write_prims_mpib_impl(False)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_mpib_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB without bundle modification and scatter under targets"""
        await self._test_read_write_prims_mpib_impl(True)

    # ----------------------------------------------------------------------
    async def _test_read_write_prims_mpib_impl(self, scatter_under_targets: bool):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB without bundle modification"""
        target_prim_paths = ["/target0", "/target1"]
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()
        stage_rt = UsdRT.Stage.Attach(usd_context.get_stage_id())

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

        cube1_prim = ogts.create_cube(stage, "Cube1", (1, 0, 0))
        cube2_prim = ogts.create_cube(stage, "Cube2", (0, 1, 0))

        (graph, [read_prims_node, _], _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.SET_VALUES: [
                    ("Write.inputs:scatterUnderTargets", scatter_under_targets),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "Write.inputs:primsBundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            targets=[cube1_prim.GetPath(), cube2_prim.GetPath()],
        )

        for target_path in target_prim_paths:
            target_xform = UsdGeom.Xform.Define(stage, target_path)
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
                target=target_xform.GetPrim().GetPath(),
            )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        # Reading the prim into a bundle again must leave the attributes intact, as we didn't modify anything
        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        graph_context = graph.get_default_graph_context()

        output_prims_bundle = graph_context.get_output_bundle(read_prims_node, "outputs_primsBundle")
        self.assertTrue(output_prims_bundle.valid)

        child_bundles = output_prims_bundle.get_child_bundles()
        self.assertEqual(len(child_bundles), 2)

        expected_colors = {
            "/Cube1": [[1.0, 0.0, 0.0]],
            "/Cube2": [[0.0, 1.0, 0.0]],
        }

        child_bundles = sort_child_bundles(child_bundles, scatter_under_targets)
        for i, bundle in enumerate(child_bundles):
            path = bundle.get_attribute_by_name("sourcePrimPath").get()

            expected_color = expected_colors[path]
            self.assertNotEqual(expected_color, None)

            # Check bundle values
            self.assertListEqual(bundle.get_attribute_by_name("primvars:displayColor").get().tolist(), expected_color)

            def check_prim_values(source_prim_path, target_prim_path, expected_color):
                # Check USD prim values
                prim = stage.GetPrimAtPath(source_prim_path)
                self.assertEqual([rgb[:] for rgb in prim.GetAttribute("primvars:displayColor").Get()], expected_color)

                # Check Fabric prim values
                prim_rt = stage_rt.GetPrimAtPath(source_prim_path)
                self.assertTrue(prim_rt.IsValid())
                self.assertEqual(
                    [rgb[:] for rgb in prim_rt.GetAttribute("primvars:displayColor").Get()], expected_color
                )

                # Check target Fabric prim values
                prim_rt = stage_rt.GetPrimAtPath(target_prim_path)
                self.assertTrue(prim_rt.IsValid())
                self.assertEqual(
                    [rgb[:] for rgb in prim_rt.GetAttribute("primvars:displayColor").Get()], expected_color
                )

            if scatter_under_targets:
                for target_prim_path in target_prim_paths:
                    check_prim_values(path, target_prim_path + path, expected_color)
            else:
                check_prim_values(path, target_prim_paths[i], expected_color)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_spib(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims SPiB without bundle modification and output to targets"""
        await self._test_read_write_prims_spib_impl(False)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_spib_no_source_prim_info(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims SPiB without bundle modification and output to targets"""
        await self._test_read_write_prims_spib_impl(False, True)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_spib_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims SPiB without bundle modification and scatter under targets"""
        await self._test_read_write_prims_spib_impl(True)

    # ----------------------------------------------------------------------
    async def _test_read_write_prims_spib_impl(
        self, scatter_under_targets: bool, remove_source_prim_info: bool = False
    ):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims SPiB without bundle modification"""
        target_prim_paths = ["/target0", "/target1"]
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()
        stage_rt = UsdRT.Stage.Attach(usd_context.get_stage_id())

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

        cube1_prim = ogts.create_cube(stage, "Cube1", (1, 0, 0))

        (graph, [read_prims_node, *_], _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("ExtractCube1", "omni.graph.nodes.ExtractPrim"),
                    ("RemoveAttrs", "omni.graph.nodes.RemoveAttribute"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.SET_VALUES: [
                    ("ExtractCube1.inputs:primPath", "/Cube1"),
                    ("RemoveAttrs.inputs:allowRemovePrimInternal", True),
                    (
                        "RemoveAttrs.inputs:attrNamesToRemove",
                        "sourcePrimPath;sourcePrimType" if remove_source_prim_info else "",
                    ),
                    ("Write.inputs:scatterUnderTargets", scatter_under_targets),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "ExtractCube1.inputs:prims"),
                    ("ExtractCube1.outputs_primBundle", "RemoveAttrs.inputs:data"),
                    ("RemoveAttrs.outputs_data", "Write.inputs:primsBundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            target=cube1_prim.GetPath(),
        )

        for target_path in target_prim_paths:
            target_xform = UsdGeom.Xform.Define(stage, target_path)
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
                target=target_xform.GetPrim().GetPath(),
            )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        # Reading the prim into a bundle again must leave the attributes intact, as we didn't modify anything
        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        graph_context = graph.get_default_graph_context()

        output_prims_bundle = graph_context.get_output_bundle(read_prims_node, "outputs_primsBundle")
        self.assertTrue(output_prims_bundle.valid)

        child_bundles = output_prims_bundle.get_child_bundles()
        self.assertEqual(len(child_bundles), 1)

        expected_colors = {
            "/Cube1": [[1.0, 0.0, 0.0]],
        }

        child_bundles = sort_child_bundles(child_bundles, scatter_under_targets)
        for i, bundle in enumerate(child_bundles):
            path = bundle.get_attribute_by_name("sourcePrimPath").get()

            expected_color = expected_colors[path]
            self.assertNotEqual(expected_color, None)

            # Check bundle values
            self.assertListEqual(bundle.get_attribute_by_name("primvars:displayColor").get().tolist(), expected_color)

            def check_prim_values(target_prim_path, expected_color):
                # Check Target USD prim values
                prim = stage.GetPrimAtPath(target_prim_path)
                # if source prim info is removed, the target prim type should not change. Otherwise it changes from Xform to Cube
                self.assertEqual(prim.GetTypeName(), "Xform" if remove_source_prim_info else "Cube")
                self.assertEqual([rgb[:] for rgb in prim.GetAttribute("primvars:displayColor").Get()], expected_color)

                # Check Target Fabric prim values
                prim_rt = stage_rt.GetPrimAtPath(target_prim_path)
                self.assertEqual(
                    [rgb[:] for rgb in prim_rt.GetAttribute("primvars:displayColor").Get()], expected_color
                )

            if scatter_under_targets:
                for target_prim_path in target_prim_paths:
                    check_prim_values(target_prim_path + path, expected_color)
            else:
                check_prim_values(target_prim_paths[i], expected_color)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_int_to_fabric(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for an integer attribute, Fabric only"""
        await self._read_write_prims_int(False, False)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_int_to_fabric_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for an integer attribute, Fabric only"""
        await self._read_write_prims_int(False, True)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_int_to_usd(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for an integer attribute, writing back to USD"""
        await self._read_write_prims_int(True, False)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_int_to_usd_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for an integer attribute, writing back to USD"""
        await self._read_write_prims_int(True, True)

    # ----------------------------------------------------------------------
    async def _read_write_prims_int(self, usd_write_back: bool, scatter_under_targets: bool):
        target_prim_paths = ["/target0", "/target1"]
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()
        stage_rt = UsdRT.Stage.Attach(usd_context.get_stage_id())

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

        cube1_prim = ogts.create_cube(stage, "Cube1", (1, 0, 0))
        cube2_prim = ogts.create_cube(stage, "Cube2", (0, 1, 0))

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("ExtractCube1", "omni.graph.nodes.ExtractPrim"),
                    ("ExtractCube2", "omni.graph.nodes.ExtractPrim"),
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                    ("SizeCube1", "omni.graph.nodes.ConstantDouble"),
                    ("SizeCube2", "omni.graph.nodes.ConstantDouble"),
                    ("InsertSizeCube1", "omni.graph.nodes.InsertAttribute"),
                    ("InsertSizeCube2", "omni.graph.nodes.InsertAttribute"),
                ],
                keys.SET_VALUES: [
                    ("ExtractCube1.inputs:primPath", "/Cube1"),
                    ("ExtractCube2.inputs:primPath", "/Cube2"),
                    ("SizeCube1.inputs:value", 200),
                    ("SizeCube2.inputs:value", 300),
                    ("InsertSizeCube1.inputs:outputAttrName", "size"),
                    ("InsertSizeCube2.inputs:outputAttrName", "size"),
                    ("Write.inputs:usdWriteBack", usd_write_back),
                    ("Write.inputs:scatterUnderTargets", scatter_under_targets),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "ExtractCube1.inputs:prims"),
                    ("Read.outputs_primsBundle", "ExtractCube2.inputs:prims"),
                    ("SizeCube1.inputs:value", "InsertSizeCube1.inputs:attrToInsert"),
                    ("SizeCube2.inputs:value", "InsertSizeCube2.inputs:attrToInsert"),
                    ("ExtractCube1.outputs_primBundle", "InsertSizeCube1.inputs:data"),
                    ("ExtractCube2.outputs_primBundle", "InsertSizeCube2.inputs:data"),
                    ("InsertSizeCube1.outputs_data", "Write.inputs:primsBundle"),
                    ("InsertSizeCube2.outputs_data", "Write.inputs:primsBundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            targets=[cube1_prim.GetPath(), cube2_prim.GetPath()],
        )

        for target_path in target_prim_paths:
            target_xform = UsdGeom.Xform.Define(stage, target_path)
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
                target=target_xform.GetPrim().GetPath(),
            )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        def check_prim_values(cube_1_path, cube_2_path):
            if usd_write_back:
                # Check USD write-back
                cube1_prim_target = stage.GetPrimAtPath(cube_1_path)
                cube2_prim_target = stage.GetPrimAtPath(cube_2_path)

                self.assertEqual(cube1_prim_target.GetAttribute("size").Get(), 200 if usd_write_back else 1)
                self.assertEqual(cube2_prim_target.GetAttribute("size").Get(), 300 if usd_write_back else 1)

                # Make sure the prim-bundle internal attributes are not written to the USD prim
                for attr_name in ["sourcePrimPath", "sourcePrimType"]:
                    self.assertEqual(cube1_prim_target.HasAttribute(attr_name), False)
                    self.assertEqual(cube2_prim_target.HasAttribute(attr_name), False)

            # Check Fabric values
            self.assertEqual(stage_rt.GetPrimAtPath(cube_1_path).GetAttribute("size").Get(), 200)
            self.assertEqual(stage_rt.GetPrimAtPath(cube_2_path).GetAttribute("size").Get(), 300)

        if scatter_under_targets:
            for target_prim_path in target_prim_paths:
                check_prim_values(target_prim_path + "/Cube1", target_prim_path + "/Cube2")
        else:
            check_prim_values(target_prim_paths[0], target_prim_paths[1])

    # ----------------------------------------------------------------------
    async def test_read_write_prims_translate_only(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an attribute pattern to just translate"""
        await self._test_read_write_prims_patterns(
            "xformOp:trans*",
            None,
            None,
            [
                [(1.0, 2.0, 3.0), (0.0, 0.0, 0.0)],
                [(4.0, 5.0, 6.0), (0.0, 0.0, 0.0)],
            ],
            False,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_translate_only_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an attribute pattern to just translate"""
        await self._test_read_write_prims_patterns(
            "xformOp:trans*",
            None,
            None,
            [
                [(1.0, 2.0, 3.0), (0.0, 0.0, 0.0)],
                [(4.0, 5.0, 6.0), (0.0, 0.0, 0.0)],
            ],
            True,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_rotate_only(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an attribute pattern to just rotate"""
        await self._test_read_write_prims_patterns(
            "xformOp:rotate*",
            None,
            None,
            [
                [(0.0, 0.0, 0.0), (7.0, 8.0, 9.0)],
                [(0.0, 0.0, 0.0), (10.0, 11.0, 12.0)],
            ],
            False,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_rotate_only_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an attribute pattern to just rotate"""
        await self._test_read_write_prims_patterns(
            "xformOp:rotate*",
            None,
            None,
            [
                [(0.0, 0.0, 0.0), (7.0, 8.0, 9.0)],
                [(0.0, 0.0, 0.0), (10.0, 11.0, 12.0)],
            ],
            True,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_xform_only(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an attribute pattern to translate and rotate"""
        await self._test_read_write_prims_patterns(
            "xformOp:*",
            None,
            None,
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            False,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_xform_only_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an attribute pattern to translate and rotate"""
        await self._test_read_write_prims_patterns(
            "xformOp:*",
            None,
            None,
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            True,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_shape1_only(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an path pattern to limit writing to Shape1 only"""
        await self._test_read_write_prims_patterns(
            None,
            "/Shape1",
            None,
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)],
            ],
            False,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_shape1_only_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an path pattern to limit writing to Shape1 only"""
        await self._test_read_write_prims_patterns(
            None,
            "/Shape1",
            None,
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)],
            ],
            True,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_shape2_only(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an path pattern to limit writing to Shape2 only"""
        await self._test_read_write_prims_patterns(
            None,
            "/Shape2",
            None,
            [
                [(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            False,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_shape2_only_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an path pattern to limit writing to Shape2 only"""
        await self._test_read_write_prims_patterns(
            None,
            "/Shape2",
            None,
            [
                [(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            True,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_both_shapes(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an path pattern to write to both shapes"""
        await self._test_read_write_prims_patterns(
            None,
            "/Shape*",
            None,
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            False,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_both_shapes_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an path pattern to write to both shapes"""
        await self._test_read_write_prims_patterns(
            None,
            "/Shape*",
            None,
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            True,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_cube_only(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an path pattern to limit writing to cubes only"""
        await self._test_read_write_prims_patterns(
            None,
            None,
            "Cube",
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)],
            ],
            False,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_cube_only_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an path pattern to limit writing to cubes only"""
        await self._test_read_write_prims_patterns(
            None,
            None,
            "Cube",
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)],
            ],
            True,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_cone_only(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an type pattern to limit writing to cones only"""
        await self._test_read_write_prims_patterns(
            None,
            None,
            "Cone",
            [
                [(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            False,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_cone_only_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an type pattern to limit writing to cones only"""
        await self._test_read_write_prims_patterns(
            None,
            None,
            "Cone",
            [
                [(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            True,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_cube_and_cone(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an type pattern to write to both cubes and cones"""
        await self._test_read_write_prims_patterns(
            None,
            None,
            "Cone;Cube",
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            False,
        )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_cube_and_cone_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims using an type pattern to write to both cubes and cones"""
        await self._test_read_write_prims_patterns(
            None,
            None,
            "Cone;Cube",
            [
                [(1.0, 2.0, 3.0), (7.0, 8.0, 9.0)],
                [(4.0, 5.0, 6.0), (10.0, 11.0, 12.0)],
            ],
            True,
        )

    # ----------------------------------------------------------------------
    async def _test_read_write_prims_patterns(
        self, attr_pattern, path_pattern, type_pattern, expected_prim_values, scatter_under_targets: bool
    ):
        target_prim_paths = ["/target0", "/target1"]

        # test USD write back since it's the main workflow currently
        usd_write_back = True

        usd_context = omni.usd.get_context()

        stage: Usd.Stage = usd_context.get_stage()
        stage_rt = UsdRT.Stage.Attach(usd_context.get_stage_id())

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

        cube_prim = ogts.create_cube(stage, "Shape1", (1, 0, 0))
        cone_prim = ogts.create_cone(stage, "Shape2", (0, 1, 0))

        # add transforms with double precision
        xform_precision = UsdGeom.XformOp.PrecisionDouble
        for prim in [cube_prim, cone_prim]:
            UsdGeom.Xformable(prim).AddTranslateOp(xform_precision).Set((0, 0, 0))
            UsdGeom.Xformable(prim).AddRotateXYZOp(xform_precision).Set((0, 0, 0))
            UsdGeom.Xformable(prim).AddScaleOp(xform_precision).Set((1, 1, 1))

        set_values = [
            ("ExtractCube.inputs:primPath", "/Shape1"),
            ("ExtractCone.inputs:primPath", "/Shape2"),
            ("TranslateCube.inputs:value", (1, 2, 3)),
            ("TranslateCone.inputs:value", (4, 5, 6)),
            ("RotateCube.inputs:value", (7, 8, 9)),
            ("RotateCone.inputs:value", (10, 11, 12)),
            ("InsertTranslateCube.inputs:outputAttrName", "xformOp:translate"),
            ("InsertTranslateCone.inputs:outputAttrName", "xformOp:translate"),
            ("InsertRotateCube.inputs:outputAttrName", "xformOp:rotateXYZ"),
            ("InsertRotateCone.inputs:outputAttrName", "xformOp:rotateXYZ"),
            ("Write.inputs:usdWriteBack", usd_write_back),
            ("Write.inputs:scatterUnderTargets", scatter_under_targets),
        ]

        if attr_pattern:
            set_values.append(("Write.inputs:attrNamesToExport", attr_pattern))

        if path_pattern:
            set_values.append(("Write.inputs:pathPattern", path_pattern))

        if type_pattern:
            set_values.append(("Write.inputs:typePattern", type_pattern))

        (graph, [extract_cube, extract_cone, *_], _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("ExtractCube", "omni.graph.nodes.ExtractPrim"),
                    ("ExtractCone", "omni.graph.nodes.ExtractPrim"),
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                    ("TranslateCube", "omni.graph.nodes.ConstantDouble3"),
                    ("TranslateCone", "omni.graph.nodes.ConstantDouble3"),
                    ("RotateCube", "omni.graph.nodes.ConstantDouble3"),
                    ("RotateCone", "omni.graph.nodes.ConstantDouble3"),
                    ("InsertTranslateCube", "omni.graph.nodes.InsertAttribute"),
                    ("InsertTranslateCone", "omni.graph.nodes.InsertAttribute"),
                    ("InsertRotateCube", "omni.graph.nodes.InsertAttribute"),
                    ("InsertRotateCone", "omni.graph.nodes.InsertAttribute"),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "ExtractCube.inputs:prims"),
                    ("Read.outputs_primsBundle", "ExtractCone.inputs:prims"),
                    ("TranslateCube.inputs:value", "InsertTranslateCube.inputs:attrToInsert"),
                    ("TranslateCone.inputs:value", "InsertTranslateCone.inputs:attrToInsert"),
                    ("RotateCube.inputs:value", "InsertRotateCube.inputs:attrToInsert"),
                    ("RotateCone.inputs:value", "InsertRotateCone.inputs:attrToInsert"),
                    ("ExtractCube.outputs_primBundle", "InsertTranslateCube.inputs:data"),
                    ("ExtractCone.outputs_primBundle", "InsertTranslateCone.inputs:data"),
                    ("InsertTranslateCube.outputs_data", "InsertRotateCube.inputs:data"),
                    ("InsertTranslateCone.outputs_data", "InsertRotateCone.inputs:data"),
                    ("InsertRotateCube.outputs_data", "Write.inputs:primsBundle"),
                    ("InsertRotateCone.outputs_data", "Write.inputs:primsBundle"),
                ],
                keys.SET_VALUES: set_values,
            },
        )

        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            targets=[cube_prim.GetPath(), cone_prim.GetPath()],
        )

        for target_path in target_prim_paths:
            target_xform = UsdGeom.Xform.Define(stage, target_path)
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
                target=target_xform.GetPrim().GetPath(),
            )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        graph_context = graph.get_default_graph_context()

        cube_bundle = graph_context.get_output_bundle(extract_cube, "outputs_primBundle")
        cone_bundle = graph_context.get_output_bundle(extract_cone, "outputs_primBundle")

        prim_bundles = [cube_bundle, cone_bundle]
        prim_paths = ["/Shape1", "/Shape2"]
        prim_types = ["Cube", "Cone"]

        # We will check these attributes
        attr_names = ["xformOp:translate", "xformOp:rotateXYZ"]

        # We need two lists of values, one for each prim
        self.assertEqual(len(expected_prim_values), 2)

        attr_pattern_tokens = attr_pattern.split(";") if attr_pattern else []
        path_pattern_tokens = path_pattern.split(";") if path_pattern else []
        type_pattern_tokens = type_pattern.split(";") if type_pattern else []

        for prim_index, _prim_bundle in enumerate(prim_bundles):
            expected_values = expected_prim_values[prim_index]

            # All attribute values must be passed
            self.assertEqual(len(expected_values), len(attr_names))

            # Check attributes values in prim bundle
            for attr_index, expected_value in enumerate(expected_values):
                attr_name = attr_names[attr_index]

                def check_prim_values(source_prim_path, target_prim_path, prim_type, attr_name, expected_value):
                    if usd_write_back:
                        prim = stage.GetPrimAtPath(target_prim_path)
                    else:
                        prim = stage_rt.GetPrimAtPath(target_prim_path)

                    # to emulate matcher behavior
                    attr_matched = True
                    if attr_pattern_tokens:
                        attr_matched = False
                        for token in attr_pattern_tokens:
                            attr_matched |= fnmatch.fnmatch(attr_name, token)
                            if attr_matched:
                                break

                    prim_matched = True
                    if path_pattern_tokens:
                        path_matched = False
                        for token in path_pattern_tokens:
                            path_matched |= fnmatch.fnmatch(source_prim_path, token)
                        prim_matched &= path_matched
                    if type_pattern_tokens:
                        type_matched = False
                        for token in type_pattern_tokens:
                            type_matched |= fnmatch.fnmatch(prim_type, token)
                        prim_matched &= type_matched

                    # Check values
                    if prim_matched:
                        if usd_write_back:
                            # usdrt syncs usd prim type if only change the type in fabric
                            self.assertEqual(prim.GetTypeName(), prim_type)
                        if attr_matched:
                            actual_value_rt = tuple(prim.GetAttribute(attr_name).Get())
                            self.assertEqual(actual_value_rt, expected_value)
                        else:
                            # if attribute does not match, it won't be exported
                            self.assertFalse(prim.GetAttribute(attr_name).IsValid())
                    else:
                        if scatter_under_targets:
                            # if unmatched, the target prim won't exist if filter failed
                            self.assertFalse(prim.IsValid())
                        else:
                            # if unmatched, the prim type won't change from xform to cube/cone
                            self.assertTrue(prim.IsValid())
                            self.assertEqual(prim.GetTypeName(), "Xform")

                if scatter_under_targets:
                    for target_prim_path in target_prim_paths:
                        check_prim_values(
                            prim_paths[prim_index],
                            target_prim_path + prim_paths[prim_index],
                            prim_types[prim_index],
                            attr_name,
                            expected_value,
                        )
                else:
                    check_prim_values(
                        prim_paths[prim_index],
                        target_prim_paths[prim_index],
                        prim_types[prim_index],
                        attr_name,
                        expected_value,
                    )

    # ----------------------------------------------------------------------
    async def test_read_write_prims_array_new_attr_to_fabric(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for a new array attribute, Fabric only"""
        await self._test_read_write_prims_vec3_array(False, False, False)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_array_new_attr_to_fabric_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for a new array attribute, Fabric only"""
        await self._test_read_write_prims_vec3_array(False, False, True)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_array_new_attr_to_usd(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for a new array attribute, writing back to USD"""
        await self._test_read_write_prims_vec3_array(True, False, False)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_array_new_attr_to_usd_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for a new array attribute, writing back to USD"""
        await self._test_read_write_prims_vec3_array(True, False, True)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_array_attr_to_fabric(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for an existing array attribute, Fabric only"""
        await self._test_read_write_prims_vec3_array(False, True, False)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_array_attr_to_fabric_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for an existing array attribute, Fabric only"""
        await self._test_read_write_prims_vec3_array(False, True, True)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_array_attr_to_usd(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for an existing array attribute, writing back to USD"""
        await self._test_read_write_prims_vec3_array(True, True, False)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_array_attr_to_usd_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims for an existing array attribute, writing back to USD"""
        await self._test_read_write_prims_vec3_array(True, True, True)

    # ----------------------------------------------------------------------
    async def _test_read_write_prims_vec3_array(
        self, usd_write_back: bool, add_new_attr: bool, scatter_under_targets: bool
    ):
        target_prim_paths = ["/target0", "/target1"]
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        stage_rt = UsdRT.Stage.Attach(usd_context.get_stage_id())

        cube_prim = ogts.create_cube(stage, "Cube", (1, 0, 0))

        attr_name = "new_array_attr"
        if not add_new_attr:
            # if graph is not going to add a new attr, make the attr now with USD API so it exists before graph execution.
            # Cannot reuse primvars:displayColor because it's a color3f[], not a strict float3[], and ConstructArray node
            # cannot set role.
            cube_prim.CreateAttribute(attr_name, Sdf.ValueTypeNames.Float3Array).Set([(1, 0, 0)])

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

        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                    ("ExtractCube", "omni.graph.nodes.ExtractPrim"),
                    ("MakeArray", "omni.graph.nodes.ConstructArray"),
                    ("InsertAttribute", "omni.graph.nodes.InsertAttribute"),
                ],
                keys.SET_VALUES: [
                    ("ExtractCube.inputs:primPath", "/Cube"),
                    ("MakeArray.inputs:arraySize", 1),
                    ("MakeArray.inputs:arrayType", "float[3][]"),
                    ("MakeArray.inputs:input0", [0.25, 0.50, 0.75]),
                    ("InsertAttribute.inputs:outputAttrName", attr_name),
                    ("Write.inputs:usdWriteBack", usd_write_back),
                    ("Write.inputs:scatterUnderTargets", scatter_under_targets),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "ExtractCube.inputs:prims"),
                    ("ExtractCube.outputs_primBundle", "InsertAttribute.inputs:data"),
                    ("MakeArray.outputs:array", "InsertAttribute.inputs:attrToInsert"),
                    ("InsertAttribute.outputs_data", "Write.inputs:primsBundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            target=cube_prim.GetPath(),
        )

        for target_path in target_prim_paths:
            target_xform = UsdGeom.Xform.Define(stage, target_path)
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
                target=target_xform.GetPrim().GetPath(),
            )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        graph_context = graph.get_default_graph_context()

        cube_bundle = graph_context.get_output_bundle(nodes[2], "outputs_primBundle")
        self.assertTrue(cube_bundle.valid)

        cube_path = "/Cube"

        def check_prim_values(source_prim_path, target_prim_path):
            # Check Fabric values
            fabric_values = stage_rt.GetPrimAtPath(target_prim_path).GetAttribute(attr_name).Get()
            self.assertEqual(len(fabric_values), 1)
            self.assertEqual(tuple(fabric_values[0]), (0.25, 0.50, 0.75))

            if usd_write_back:
                # With USD write back, the existing attribute should be updated in USD
                usd_values = stage.GetPrimAtPath(target_prim_path).GetAttribute(attr_name).Get()
                self.assertEqual(len(usd_values), 1)
                self.assertEqual(tuple(usd_values[0]), (0.25, 0.50, 0.75))
            else:
                if scatter_under_targets:
                    self.assertFalse(stage.GetPrimAtPath(target_prim_path).IsValid())

        if scatter_under_targets:
            for target_prim_path in target_prim_paths:
                check_prim_values(cube_path, target_prim_path + cube_path)
        else:
            # only one source
            check_prim_values(cube_path, target_prim_paths[0])

    async def test_read_write_prims_remove_missing_to_usd(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB with remove missing attributes and prims matching to USD"""
        await self._test_read_write_prims_remove_missing(True, False)

    async def test_read_write_prims_remove_missing_to_usd_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB with remove missing attributes and prims matching to USD"""
        await self._test_read_write_prims_remove_missing(True, True)

    async def test_read_write_prims_remove_missing_to_fabric(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB with remove missing attributes and prims matching to Fabric"""
        await self._test_read_write_prims_remove_missing(False, False)

    async def test_read_write_prims_remove_missing_to_fabric_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB with remove missing attributes and prims matching to Fabric"""
        await self._test_read_write_prims_remove_missing(False, True)

    # ----------------------------------------------------------------------
    async def _test_read_write_prims_remove_missing(self, usd_write_back: bool, scatter_under_targets: bool):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB with exact attributes and prims matching"""
        usd_context = omni.usd.get_context()
        stage_usd: Usd.Stage = usd_context.get_stage()
        stage_rt = UsdRT.Stage.Attach(usd_context.get_stage_id())

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

        # Stage hierarchy:
        #
        # /Xform
        #   /Xform
        #     /Cube2
        #   /Cube1
        # /target0
        # /target1

        stage_usd.DefinePrim("/Xform", "Xform")
        stage_usd.DefinePrim("/Xform/Xform", "Xform")

        cube1_prim = stage_usd.DefinePrim("/Xform/Cube1", "Cube")
        cube2_prim = stage_usd.DefinePrim("/Xform/Xform/Cube2", "Cube")

        test_attr_name = "test_attr"
        test_target_existing_attr_name = "test_attr2"
        test_both_existing_attr_name = "test_attr3"

        cube1_prim.CreateAttribute(test_attr_name, Sdf.ValueTypeNames.Float3)
        cube2_prim.CreateAttribute(test_attr_name, Sdf.ValueTypeNames.Float3)

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.SET_VALUES: [
                    ("Write.inputs:usdWriteBack", usd_write_back),
                    ("Write.inputs:scatterUnderTargets", scatter_under_targets),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "Write.inputs:primsBundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage_usd.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            targets=[cube1_prim.GetPath(), cube2_prim.GetPath()],
        )

        target_prim_paths = ["/target0", "/target1"]
        for target_path in target_prim_paths:
            target_xform = UsdGeom.Cube.Define(stage_usd, target_path)
            omni.kit.commands.execute(
                "AddRelationshipTarget",
                relationship=stage_usd.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
                target=target_xform.GetPrim().GetPath(),
            )

            # put a Sphere under the target. This should NOT be removed by removeMissingPrims later, as it doesn't exist in the upstream bundle
            stage_usd.DefinePrim(target_path + "/Sphere", "Sphere")

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        # check attributes are exported correctly and setup the test scenario for removeMissingAttrs
        # Add test_target_existing_attr_name to target prims only
        # Add test_both_existing_attr_name to both source and target prims
        for i, path in enumerate([cube1_prim.GetPath(), cube2_prim.GetPath()]):

            def check_initial_prim(target_prim_path):
                for stage in [stage_rt, stage_usd]:
                    # Check prim
                    prim = stage.GetPrimAtPath(target_prim_path)
                    self.assertTrue(prim.IsValid())

                    # Check attribute
                    attr = prim.GetAttribute(test_attr_name)
                    self.assertTrue(attr.IsValid())

                    # skip stage_usd check if usd_write_back is off
                    if not usd_write_back:
                        break

            def add_attribute_to_source(source_prim_path):
                # creates additional attributes on output prims.
                # these attribute should not be removed, as there did not exist in upstream bundle
                for stage, sdf in [(stage_rt, SdfRT), (stage_usd, Sdf)]:
                    prim = stage.GetPrimAtPath(source_prim_path)
                    prim.CreateAttribute(test_both_existing_attr_name, sdf.ValueTypeNames.Float2, True)

                    if not usd_write_back:
                        break

            def add_attribute_to_target_output(target_prim_path):
                # creates additional attributes on output prims.
                # these attribute should not be removed, as there did not exist in upstream bundle
                for stage, sdf in [(stage_rt, SdfRT), (stage_usd, Sdf)]:
                    prim = stage.GetPrimAtPath(target_prim_path)
                    prim.CreateAttribute(test_target_existing_attr_name, sdf.ValueTypeNames.Float2, True)
                    prim.CreateAttribute(test_both_existing_attr_name, sdf.ValueTypeNames.Float2, True)

                    if not usd_write_back:
                        break

            if scatter_under_targets:
                for target_prim_path in target_prim_paths:
                    check_initial_prim(target_prim_path + path.pathString)
                    add_attribute_to_target_output(target_prim_path + path.pathString)
            else:
                check_initial_prim(target_prim_paths[i])
                add_attribute_to_target_output(target_prim_paths[i])

            add_attribute_to_source(path.pathString)

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        # --------------------------- test removeMissingAttrs -----------------------------------------
        # remove the test_attr from source and tick the graph, it must be removed from target too
        cube1_prim.RemoveProperty(test_attr_name)
        cube2_prim.RemoveProperty(test_attr_name)

        # remove the test_attr from source and tick the graph, it must be removed from target too
        cube1_prim.RemoveProperty(test_both_existing_attr_name)
        cube2_prim.RemoveProperty(test_both_existing_attr_name)

        # TODO ReadPrims node won't clear the attributes from output bundle due to lack of USD change tracking,
        # has to set attrNamesToImport to force a re-evaluation
        # FIXME when ReadPrims node properly tracks USD changes.
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: [
                    # changing filter from '*' to '**' applies same filter but forces node to re-evaluate to remove deleted attributes
                    ("Read.inputs:attrNamesToImport", "**"),
                ],
            },
        )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        for i, path in enumerate([cube1_prim.GetPath(), cube2_prim.GetPath()]):

            def check_prim_attr_matched(source_prim_path, target_prim_path):
                for stage in [stage_rt, stage_usd]:
                    prim = stage.GetPrimAtPath(target_prim_path)
                    self.assertTrue(prim.IsValid())

                    # Check attribute, it should be gone
                    attr = prim.GetAttribute(test_attr_name)
                    self.assertFalse(attr.IsValid())

                    # Check pre-existing attribute that has not been written to, it should stay
                    attr = prim.GetAttribute(test_target_existing_attr_name)
                    self.assertTrue(attr.IsValid())

                    # Check pre-existing attribute that has been written to, it should stay
                    attr = prim.GetAttribute(test_both_existing_attr_name)
                    self.assertTrue(attr.IsValid())

                    # skip stage_usd check if usd_write_back is off
                    if not usd_write_back:
                        break

            if scatter_under_targets:
                for target_prim_path in target_prim_paths:
                    check_prim_attr_matched(path, target_prim_path + path.pathString)
            else:
                check_prim_attr_matched(path, target_prim_paths[i])

        # ------------------------- test removeMissingPrims ---------------------------------
        # remove the /Cube1 from source and tick the graph, it must be removed from target too
        omni.kit.commands.execute(
            "RemoveRelationshipTarget",
            relationship=stage_usd.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            target=cube2_prim.GetPath(),
        )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        def check_prim_matched(target_root_path, prim_path, valid: bool):
            # Check USD and fabric prim, "Cube2" should be removed, "Cube1" stays
            target_prim_path = target_root_path + prim_path
            for stage in [stage_rt, stage_usd]:
                prim = stage.GetPrimAtPath(target_prim_path)
                self.assertEqual(prim.IsValid(), valid)

                # The pre-existing Sphere prim should still exist
                self.assertTrue(stage.GetPrimAtPath(target_root_path + "/Sphere").IsValid())

                # The stub prim /targetN/Xform/Xform should also be removed
                self.assertFalse(stage.GetPrimAtPath(target_root_path + "/Xform/Xform").IsValid())

                # skip stage_usd check if usd_write_back is off
                if not usd_write_back:
                    break

        if scatter_under_targets:
            for i, path in enumerate([cube1_prim.GetPath(), cube2_prim.GetPath()]):
                for target_prim_path in target_prim_paths:
                    check_prim_matched(target_prim_path, path.pathString, i == 0)
        else:
            check_prim_matched(target_prim_paths[0], "", True)

        # test removeMissingPrims with no prim input
        # remove all prims from ReadPrims node, and the target should be empty
        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage_usd.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            targets=[],
        )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        # target prims should have only the "Sphere" prim as child
        for target_prim_path in target_prim_paths:
            target_prim = stage_usd.GetPrimAtPath(target_prim_path)
            self.assertEqual(target_prim.GetAllChildrenNames(), ["Sphere"])

    async def test_read_write_prims_remove_schema_properties(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB that removes attributes defined within schema"""
        await self._test_read_write_prims_remove_schema_properties(False)

    async def test_read_write_prims_remove_schema_properties_scatter(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB that removes attributes defined within schema"""
        await self._test_read_write_prims_remove_schema_properties(True)

    # ----------------------------------------------------------------------
    async def _test_read_write_prims_remove_schema_properties(self, scatter_under_targets: bool):
        """
        Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB that removes attributes defined within schema
        This is a USD-only test as Fabric doesn't have the schema concept yet
        """
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()

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

        # Stage hierarchy:
        #
        # /Xform
        # /target

        xform_prim = stage.DefinePrim("/Xform", "Xform")

        # both are schema properties of Xformable
        test_attr_name = UsdGeom.Tokens.visibility
        test_rel_name = UsdGeom.Tokens.proxyPrim

        xform_prim.GetAttribute(UsdGeom.Tokens.visibility).Set(UsdGeom.Tokens.invisible)
        xform_prim.GetRelationship(test_rel_name).AddTarget("/Foo/bar")

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.SET_VALUES: [
                    ("Write.inputs:usdWriteBack", True),
                    ("Write.inputs:scatterUnderTargets", scatter_under_targets),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "Write.inputs:primsBundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            targets=[xform_prim.GetPath()],
        )

        target_xform = UsdGeom.Xform.Define(stage, "/target")
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
            target=target_xform.GetPrim().GetPath(),
        )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        def check_initial_prim(target_prim_path):
            # Check prim
            prim = stage.GetPrimAtPath(target_prim_path)
            self.assertTrue(prim.IsValid())

            # Check attribute
            # It should be both valid and authored for target prim
            attr = prim.GetAttribute(test_attr_name)
            self.assertTrue(attr.IsAuthored())
            self.assertEqual(attr.Get(), UsdGeom.Tokens.invisible)

            # Check relationship
            # It should be both valid and authored for target prim
            rel = prim.GetRelationship(test_rel_name)
            self.assertTrue(rel.IsAuthored())
            self.assertListEqual(rel.GetTargets(), ["/Foo/bar"])

        if scatter_under_targets:
            check_initial_prim("/target/Xform")
        else:
            check_initial_prim("/target")

        # Remove spec for both properties
        # XXX cannot use RemoveProperty due to change tracking not working properly in ReadPrimsV2 node
        # set attrNamesToImport to empty instead to clear the attributes from ReadPrimsV2 output bundle
        # xform_prim.RemoveProperty(test_attr_name)
        # xform_prim.RemoveProperty(test_rel_name)
        controller.edit(self.TEST_GRAPH_PATH, {keys.SET_VALUES: ("Read.inputs:attrNamesToImport", "")})

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        def check_ticked_prim(target_prim_path):
            # Check prim
            prim = stage.GetPrimAtPath(target_prim_path)
            self.assertTrue(prim.IsValid())

            # Check attribute
            # It should still exits but be reverted to unauthored schema value
            attr = prim.GetAttribute(test_attr_name)
            self.assertTrue(attr.IsValid())
            self.assertFalse(attr.IsAuthored())
            self.assertEqual(attr.Get(), UsdGeom.Tokens.inherited)

            # Check relationship
            # It should still exits but be reverted to unauthored schema value
            rel = prim.GetRelationship(test_rel_name)
            self.assertTrue(rel.IsValid())
            self.assertFalse(rel.IsAuthored())
            self.assertListEqual(rel.GetTargets(), [])

        if scatter_under_targets:
            check_ticked_prim("/target/Xform")
        else:
            check_ticked_prim("/target")

    # ----------------------------------------------------------------------
    async def test_read_write_prims_fill_usd_hierarchy_gap(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims when some prim along target hierarchy does not fully exist"""
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()

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

        UsdGeom.Xform.Define(stage, "/Xform")
        scope = UsdGeom.Scope.Define(stage, "/Xform/Scope")
        cube = UsdGeom.Cube.Define(stage, "/Xform/Scope/Cube")

        target_xform = UsdGeom.Xform.Define(stage, "/Target")

        (graph, [_, _], _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.SET_VALUES: [
                    ("Write.inputs:usdWriteBack", True),
                    ("Write.inputs:scatterUnderTargets", True),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "Write.inputs:primsBundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            target=cube.GetPrim().GetPath(),
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
            target=target_xform.GetPrim().GetPath(),
        )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        # check the target hierarchy was created, all missing prims has the type Xform
        self.assertTrue(UsdGeom.Xform.Get(stage, "/Target/Xform"))
        self.assertTrue(UsdGeom.Xform.Get(stage, "/Target/Xform/Scope"))
        self.assertTrue(UsdGeom.Cube.Get(stage, "/Target/Xform/Scope/Cube"))

        # add /Xform/Scope to import targets
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            target=scope.GetPrim().GetPath(),
        )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        # check the target hierarchy was created, because "/Xform/Scope" is also exported,
        # the type should be changed from Xform to Scope
        self.assertTrue(UsdGeom.Xform.Get(stage, "/Target/Xform"))
        self.assertTrue(UsdGeom.Scope.Get(stage, "/Target/Xform/Scope"))
        self.assertFalse(UsdGeom.Xform.Get(stage, "/Target/Xform/Scope"))
        self.assertTrue(UsdGeom.Cube.Get(stage, "/Target/Xform/Scope/Cube"))

    # ----------------------------------------------------------------------
    async def test_read_write_prims_find_prims_loop_guard(self):
        await self._test_read_write_prims_find_prims_loop_guard(False)

    # ----------------------------------------------------------------------
    async def test_read_write_prims_find_prims_loop_guard_scatter(self):
        await self._test_read_write_prims_find_prims_loop_guard(True)

    # ----------------------------------------------------------------------
    async def _test_read_write_prims_find_prims_loop_guard(self, scatter_under_targets: bool):
        """Test omni.graph.nodes.ReadPrimsV2 find prims tracker and WritePrims that loop can be detect and prevented"""
        logging = carb.logging.acquire_logging()

        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()

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

        xform = UsdGeom.Xform.Define(stage, "/Xform")
        xform_xform = UsdGeom.Xform.Define(stage, "/Xform/Xform")

        xform2 = UsdGeom.Xform.Define(stage, "/Xform2")
        xform2_xform = UsdGeom.Xform.Define(stage, "/Xform2/Xform")

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("Read2", "omni.graph.nodes.ReadPrimsV2"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.SET_VALUES: [
                    ("Read.inputs:pathPattern", "*"),
                    ("Write.inputs:usdWriteBack", True),
                    ("Write.inputs:scatterUnderTargets", scatter_under_targets),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "Write.inputs:primsBundle"),
                ],
            },
        )

        # ------------ test path tracker on connected Read node ------------
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            target=xform.GetPrim().GetPath(),
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
            target=xform_xform.GetPrim().GetPath(),
        )

        log_was_enabled = logging.is_log_enabled()
        # turn off carb logging so the error does not trigger test failure
        # og logging still works and logs node error message.
        logging.set_log_enabled(False)

        await self._evaluate_graph(graph, controller)
        error_message = self._get_node_messages(graph, og.ERROR)

        # check the error is logged
        self.assertEqual(
            error_message,
            [
                "[/World/TestGraph] Target prim '/Xform/Xform' is a descendant of FindPrims/ReadPrims source '/Xform', writing to the target may mutate the input state. Please choose a different target"
            ],
        )

        # check the target prim is NOT created
        if scatter_under_targets:
            self.assertFalse(stage.GetPrimAtPath("/Xform/Xform/Xform").IsValid())

        # ------------ test path tracker on unconnected Read node ------------
        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: [
                    ("Read2.inputs:pathPattern", "*"),
                ]
            },
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read2.inputs:prims"),
            target=xform2.GetPrim().GetPath(),
        )

        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
            targets=[xform2_xform.GetPrim().GetPath()],
        )

        await self._evaluate_graph(graph, controller)
        error_message = self._get_node_messages(graph, og.ERROR)

        # check the error is logged
        self.assertEqual(
            error_message,
            [
                "[/World/TestGraph] Target prim '/Xform2/Xform' is a descendant of FindPrims/ReadPrims source '/Xform2', writing to the target may mutate the input state. Please choose a different target"
            ],
        )

        # check the target prim is NOT created
        if scatter_under_targets:
            self.assertFalse(stage.GetPrimAtPath("/Xform2/Xform/Xform").IsValid())

        logging.set_log_enabled(log_was_enabled)

    # ----------------------------------------------------------------------
    async def test_write_prims_empty_bundle(self):
        """Test warnings and errors issued by omni.graph.nodes.WritePrimsV2 when its input bundle is empty"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("bundle_ctor", "omni.graph.nodes.BundleConstructor"),
                    ("write_prims", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.CONNECT: [
                    ("bundle_ctor.outputs_bundle", "write_prims.inputs:primsBundle"),
                ],
            },
        )

        with self.CarbLogSuppressor():
            await self._evaluate_graph(graph, controller)

            # no warning or errors messages are expected
            self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

    # ----------------------------------------------------------------------
    async def test_write_prims_source_prim_path_validation_str(self):
        """Test warnings and errors issued by omni.graph.nodes.WritePrimsV2 when a string sourcePrimPath is detected"""
        await self._test_write_prims_attribute_validation_impl(
            "sourcePrimPath", "foo", "String", "sourcePrimPath attribute must hold a token"
        )

    async def test_write_prims_source_prim_path_validation_int(self):
        """Test warnings and errors issued by omni.graph.nodes.WritePrimsV2 when an integer sourcePrimPath is detected"""
        await self._test_write_prims_attribute_validation_impl(
            "sourcePrimPath", 42, "Int", "sourcePrimPath attribute must hold a token"
        )

    async def test_write_prims_source_prim_type_validation_str(self):
        """Test warnings and errors issued by omni.graph.nodes.WritePrimsV2 when a string sourcePrimType is detected"""
        await self._test_write_prims_attribute_validation_impl(
            "sourcePrimType", "bar", "String", "sourcePrimType attribute must hold a token"
        )

    async def test_write_prims_source_prim_type_validation_int(self):
        """Test warnings and errors issued by omni.graph.nodes.WritePrimsV2 when an integer rsourcePrimType is detected"""
        await self._test_write_prims_attribute_validation_impl(
            "sourcePrimType", 42, "Int", "sourcePrimType attribute must hold a token"
        )

    async def _test_write_prims_attribute_validation_impl(
        self, attr_name, attr_value, attr_type, expected_error_message
    ):
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        ogts.create_cube(stage, "Cube", (1, 0, 0))

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("CreateEmptyBundle", "omni.graph.nodes.BundleConstructor"),
                    ("WritePrims", "omni.graph.nodes.WritePrimsV2"),
                    ("ConstSourcePrimPath", f"omni.graph.nodes.Constant{attr_type}"),
                    ("InsertSourcePrimPath", "omni.graph.nodes.InsertAttribute"),
                ],
                keys.SET_VALUES: [
                    ("ConstSourcePrimPath.inputs:value", attr_value),
                    ("InsertSourcePrimPath.inputs:outputAttrName", attr_name),
                ],
                keys.CONNECT: [
                    ("ConstSourcePrimPath.inputs:value", "InsertSourcePrimPath.inputs:attrToInsert"),
                    ("CreateEmptyBundle.outputs_bundle", "InsertSourcePrimPath.inputs:data"),
                    ("InsertSourcePrimPath.outputs_data", "WritePrims.inputs:primsBundle"),
                ],
            },
        )

        target_xform = UsdGeom.Xform.Define(stage, "/Targets")
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/WritePrims.inputs:prims"),
            target=target_xform.GetPrim().GetPath(),
        )

        with self.CarbLogSuppressor():
            await self._evaluate_graph(graph, controller)

            # check the error is logged
            expected_message = f"[{self.TEST_GRAPH_PATH}] Bundle '{self.TEST_GRAPH_PATH}/__runtime_data/InsertSourcePrimPath_outputs_data' was skipped because it was invalid! {expected_error_message}"
            self.assertEqual(self._get_node_messages(graph, og.WARNING), [expected_message])

            # no warning message is expected
            self._assert_no_compute_messages(graph, [og.ERROR])

    # ----------------------------------------------------------------------
    async def test_read_write_prims_to_layer(self):
        """Test omni.graph.nodes.ReadPrimsV2 and WritePrims MPiB to a different layer"""
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()

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

        session_layer = stage.GetSessionLayer()
        sub_layer = Sdf.Layer.CreateAnonymous()
        sub_layer2 = Sdf.Layer.CreateAnonymous()
        root_layer = stage.GetRootLayer()
        root_layer.subLayerPaths.append(sub_layer.identifier)
        root_layer.subLayerPaths.append(sub_layer2.identifier)

        test_layers = [session_layer, sub_layer, sub_layer2, root_layer]

        # switch edit target to sub_layer2
        with Usd.EditContext(stage, sub_layer2):

            def create_scope_with_test_attrs(prim_path: str):
                scope = UsdGeom.Scope.Define(stage, prim_path)
                scope_prim = scope.GetPrim()
                scope_prim.CreateAttribute("test_attr_0", Sdf.ValueTypeNames.Bool).Set(True)
                scope_prim.CreateAttribute("test_attr_1", Sdf.ValueTypeNames.Bool).Set(True)
                scope_prim.CreateAttribute("test_attr_2", Sdf.ValueTypeNames.Bool).Set(True)
                scope_prim.CreateAttribute("test_attr_3", Sdf.ValueTypeNames.Bool).Set(True)
                return scope_prim

            test_layer_count = len(test_layers)

            source_prims = []
            for i in range(test_layer_count):
                source_prims.append(create_scope_with_test_attrs(f"/Scope{i}"))

            # make a graph to Read from scope_prim_N, Write to scope_N_out prim with only test_attr_N to layer index N
            (graph, _, _, _) = controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.CREATE_NODES: [
                        ("Read0", "omni.graph.nodes.ReadPrimsV2"),
                        ("Read1", "omni.graph.nodes.ReadPrimsV2"),
                        ("Read2", "omni.graph.nodes.ReadPrimsV2"),
                        ("Read3", "omni.graph.nodes.ReadPrimsV2"),
                        ("Write0", "omni.graph.nodes.WritePrimsV2"),
                        ("Write1", "omni.graph.nodes.WritePrimsV2"),
                        ("Write2", "omni.graph.nodes.WritePrimsV2"),
                        ("Write3", "omni.graph.nodes.WritePrimsV2"),
                    ],
                    keys.SET_VALUES: [
                        ("Write0.inputs:layerIdentifier", "<Session Layer>"),
                        ("Write1.inputs:layerIdentifier", sub_layer.identifier),
                        ("Write2.inputs:layerIdentifier", ""),  # write to current edit target (sub_layer2 in this case)
                        ("Write3.inputs:layerIdentifier", "<Root Layer>"),
                        ("Write0.inputs:attrNamesToExport", "test_attr_0"),
                        ("Write1.inputs:attrNamesToExport", "test_attr_1"),
                        ("Write2.inputs:attrNamesToExport", "test_attr_2"),
                        ("Write3.inputs:attrNamesToExport", "test_attr_3"),
                    ],
                    keys.CONNECT: [
                        ("Read0.outputs_primsBundle", "Write0.inputs:primsBundle"),
                        ("Read1.outputs_primsBundle", "Write1.inputs:primsBundle"),
                        ("Read2.outputs_primsBundle", "Write2.inputs:primsBundle"),
                        ("Read3.outputs_primsBundle", "Write3.inputs:primsBundle"),
                    ],
                },
            )

            for i in range(test_layer_count):
                omni.kit.commands.execute(
                    "SetRelationshipTargets",
                    relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read{i}.inputs:prims"),
                    targets=[source_prims[i].GetPath()],
                )

                omni.kit.commands.execute(
                    "SetRelationshipTargets",
                    relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write{i}.inputs:prims"),
                    targets=[f"/Scope_{i}_out"],
                )

            await self._evaluate_graph(graph, controller)
            self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

            # Scope_0_out should only exist on test_layers[0] with test_attr_0
            # Scope_1_out should only exist on test_layers[1] with test_attr_1
            # Scope_2_out should only exist on test_layers[2] with test_attr_2
            # Scope_3_out should only exist on test_layers[3] with test_attr_3
            for layer_idx, layer in enumerate(test_layers):
                for prim_idx in range(test_layer_count):
                    prim_out_spec = layer.GetPrimAtPath(f"/Scope_{prim_idx}_out")
                    if prim_idx == layer_idx:
                        self.assertTrue(prim_out_spec)
                        self.assertEqual(prim_out_spec.specifier, Sdf.SpecifierDef)

                        for attr_idx in range(test_layer_count):
                            attr_out_spec = layer.GetAttributeAtPath(f"/Scope_{prim_idx}_out.test_attr_{attr_idx}")
                            if attr_idx == layer_idx:
                                self.assertTrue(attr_out_spec)
                                self.assertTrue(attr_out_spec.default)
                            else:
                                self.assertFalse(attr_out_spec)
                    else:
                        self.assertFalse(prim_out_spec)

            #############################################################
            # Change the target to session layer and check `over` spec
            test_layer_index = 2

            # change value from True to False
            source_prims[test_layer_index].GetAttribute(f"test_attr_{test_layer_index}").Set(False)

            controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.SET_VALUES: [
                        (f"Write{test_layer_index}.inputs:layerIdentifier", "<Session Layer>"),
                    ],
                },
            )

            await self._evaluate_graph(graph, controller)
            self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

            prim_out_spec = session_layer.GetPrimAtPath(f"/Scope_{test_layer_index}_out")
            self.assertTrue(prim_out_spec)
            self.assertEqual(prim_out_spec.specifier, Sdf.SpecifierOver)

            attr_out_spec = session_layer.GetAttributeAtPath(
                f"/Scope_{test_layer_index}_out.test_attr_{test_layer_index}"
            )
            self.assertTrue(attr_out_spec)
            self.assertFalse(attr_out_spec.default)  # value should be set to False because we changed it

            prev_attr_out_spec = test_layers[test_layer_index].GetAttributeAtPath(
                f"/Scope_{test_layer_index}_out.test_attr_{test_layer_index}"
            )
            self.assertTrue(prev_attr_out_spec)
            self.assertTrue(
                prev_attr_out_spec.default
            )  # the value should still be True as it should not be written to anymore

    # ----------------------------------------------------------------------
    async def test_read_write_prims_layer_identifier_resolver(self):
        """
        Test omni.graph.nodes.ReadPrimsV2 and WritePrimsV2 layer identifier resolver

        Construct file hierarchy like this:
        /Ref/ref.usda
        /Ref/ref_sublayer.usda
        /main.usda

        /Ref/ref_sublayer.usda is a sublayer of /Ref/ref.usda
        A graph with WritePrimsV2 node is in /Ref/ref.usda, with a layerIdentifier to ./ref_sublayer.usda
        If /Ref/ref.usda is opened in Kit and graph executes, the output should be written to /Ref/ref_sublayer.usda

        /Ref/ref_sublayer.usda is also a sublayer of /main.usda
        /main.usda also references /Ref/ref.usda
        If /main.usda is opened in Kit and graph executes, the relative layerIdentifier of ./ref_sublayer.usda should
        still be resolved to /Ref/ref_sublayer.usda, but not /ref_sublayer.usda, and graph out put on that layer.
        """
        with tempfile.TemporaryDirectory() as tmpdirname:
            usd_context = omni.usd.get_context()
            controller = og.Controller()
            keys = og.Controller.Keys

            ref_dir = os.path.join(tmpdirname, "Ref")

            ref_sublayer_fn = os.path.join(ref_dir, "ref_sub_layer.usda")
            ref_sublayer = Sdf.Layer.CreateNew(ref_sublayer_fn)
            ref_sublayer.Save()

            ref_fn = os.path.join(ref_dir, "ref.usda")
            Sdf.Layer.CreateNew(ref_fn).Save()

            success, error = await usd_context.open_stage_async(str(ref_fn))
            self.assertTrue(success, error)

            stage = usd_context.get_stage()
            root_layer = stage.GetRootLayer()
            root_layer.subLayerPaths.append("./ref_sub_layer.usda")

            world = UsdGeom.Xform.Define(stage, "/World")
            stage.SetDefaultPrim(world.GetPrim())

            source = UsdGeom.Scope.Define(stage, "/World/Source")
            source_prim = source.GetPrim()
            source_prim.CreateAttribute("test_attr", Sdf.ValueTypeNames.Bool).Set(True)

            (graph, _, _, _) = controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.CREATE_NODES: [
                        ("Read", "omni.graph.nodes.ReadPrimsV2"),
                        ("Write", "omni.graph.nodes.WritePrimsV2"),
                    ],
                    keys.SET_VALUES: [
                        ("Write.inputs:layerIdentifier", "./ref_sub_layer.usda"),
                        ("Write.inputs:attrNamesToExport", "test_attr"),
                    ],
                    keys.CONNECT: [
                        ("Read.outputs_primsBundle", "Write.inputs:primsBundle"),
                    ],
                },
            )

            omni.kit.commands.execute(
                "SetRelationshipTargets",
                relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
                targets=["/World/Source"],
            )

            omni.kit.commands.execute(
                "SetRelationshipTargets",
                relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
                targets=["/World/Target"],
            )

            # Block saving the stage first before graph ticking. The ticked result should not be saved and discard later
            result = usd_context.save_stage()
            self.assertTrue(result)

            # Tick the graph to make sure the layerIdentifier is resolved correctly within /Ref/ref.usda stage
            await self._evaluate_graph(graph, controller)
            self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

            prim_out_spec = ref_sublayer.GetPrimAtPath("/World/Target")
            self.assertTrue(prim_out_spec)
            self.assertEqual(prim_out_spec.specifier, Sdf.SpecifierDef)
            attr_out_spec = ref_sublayer.GetAttributeAtPath("/World/Target.test_attr")
            self.assertTrue(attr_out_spec)
            self.assertTrue(attr_out_spec.default)

            prim_out_spec = root_layer.GetPrimAtPath("/World/Target")
            self.assertFalse(prim_out_spec)

            # close the stage and release layer ref to discard all temporary changes.
            await omni.usd.get_context().close_stage_async()
            ref_sublayer = None

            # Now make the main.usda, add Ref/ref.usda as a reference and tick graph
            main_fn = os.path.join(tmpdirname, "main.usda")
            Sdf.Layer.CreateNew(main_fn).Save()

            success, error = await usd_context.open_stage_async(str(main_fn))
            self.assertTrue(success, error)
            stage = usd_context.get_stage()

            # Add "./Ref/ref_sub_layer.usda" again as a local sublayer so graph can resolve and write to it on this new stage
            root_layer = stage.GetRootLayer()
            root_layer.subLayerPaths.append("./Ref/ref_sub_layer.usda")

            # For whatever reason need to make an empty graph otherwise test crashes when referencing in "./Ref/ref.usda"
            # same code does NOT crash while executed within Kit Script Editor without the dummy graph
            controller.edit("/dummy", {})
            ref_prim = stage.DefinePrim("/Ref")
            refs = ref_prim.GetReferences()
            refs.AddReference("./Ref/ref.usda")

            # reopen Ref/ref_sublayer.usda
            ref_sublayer = Sdf.Layer.FindOrOpen(ref_sublayer_fn)
            self.assertTrue(ref_sublayer)

            # double check the target spec does not exist before tick the graph
            self.assertFalse(ref_sublayer.GetPrimAtPath("/Ref/Target"))

            # let OG populate the referenced graph
            await omni.kit.app.get_app().next_update_async()

            # Tick the graph to make sure the layerIdentifier is resolved correctly when referencing /Ref/ref.usda
            graph = og.get_graph_by_path("/Ref/TestGraph")
            await self._evaluate_graph(graph, controller)
            self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

            # the references graph should be able to resolve layerIdentifier even if the relative path anchors differently
            prim_out_spec = ref_sublayer.GetPrimAtPath("/Ref/Target")
            self.assertTrue(prim_out_spec)
            self.assertEqual(prim_out_spec.specifier, Sdf.SpecifierDef)
            attr_out_spec = ref_sublayer.GetAttributeAtPath("/Ref/Target.test_attr")
            self.assertTrue(attr_out_spec)
            self.assertTrue(attr_out_spec.default)

            prim_out_spec = root_layer.GetPrimAtPath("/Ref/Target")
            self.assertFalse(prim_out_spec)

    # ----------------------------------------------------------------------
    async def test_write_prims_relationship(self):
        """Test WritePrims MPiB to export a relationship"""
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()
        stage_rt = UsdRT.Stage.Attach(usd_context.get_stage_id())

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

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("RelBundle", "__test_write_prims_v2__.relationship_bundle"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.CONNECT: [
                    ("RelBundle.outputs_out_0", "Write.inputs:primsBundle"),
                ],
            },
        )

        target_path = "/Target"
        target_xform = UsdGeom.Xform.Define(stage, target_path)
        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
            targets=[target_xform.GetPrim().GetPath()],
        )

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        # Check USD prim relationship
        prim = stage.GetPrimAtPath(target_path)
        rel = prim.GetRelationship(REL_BUNDLE_TRL_NAME)
        self.assertTrue(rel)
        self.assertListEqual(REL_BUNDLE_TARGETS, rel.GetTargets())

        # Check Fabric prim relationship
        prim_rt = stage_rt.GetPrimAtPath(target_path)
        rel_rt = prim_rt.GetRelationship(REL_BUNDLE_TRL_NAME)
        self.assertTrue(rel_rt)
        self.assertListEqual(REL_BUNDLE_TARGETS_RT, rel_rt.GetTargets())

    # ----------------------------------------------------------------------
    async def test_write_prims_interpolation(self):
        """Test writing interpolation metadata through omni.graph.nodes.WritePrimsV2"""
        target_prim_path = "/target"
        usd_context = omni.usd.get_context()
        stage: Usd.Stage = usd_context.get_stage()

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

        xform_prim = UsdGeom.Xform.Define(stage, target_prim_path)
        cube_prim = ogts.create_cube(stage, "Cube1", (1, 0, 0))

        (graph, [_, _], _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("Read", "omni.graph.nodes.ReadPrimsV2"),
                    ("Write", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.SET_VALUES: [
                    ("Write.inputs:scatterUnderTargets", True),
                ],
                keys.CONNECT: [
                    ("Read.outputs_primsBundle", "Write.inputs:primsBundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Read.inputs:prims"),
            targets=[cube_prim.GetPath()],
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/Write.inputs:prims"),
            target=xform_prim.GetPrim().GetPath(),
        )

        # Set the interpolation on a primvar
        interpolation = "uniform"
        primvar = UsdGeom.Primvar(cube_prim.GetAttribute("primvars:displayColor"))
        self.assertTrue(primvar)
        primvar.SetInterpolation(interpolation)

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        target_prim_path = xform_prim.GetPath().AppendPath(
            cube_prim.GetPath().MakeRelativePath(Sdf.Path.absoluteRootPath)
        )
        target_prim = stage.GetPrimAtPath(target_prim_path)
        self.assertTrue(target_prim)
        attr = target_prim.GetAttribute("primvars:displayColor")
        self.assertTrue(attr)
        primvar = UsdGeom.Primvar(attr)
        self.assertTrue(primvar)
        self.assertEqual(primvar.GetInterpolation(), interpolation)

        # Change the interpolation of the input prim, then evaluate the graph again...
        interpolation = "constant"
        primvar = UsdGeom.Primvar(cube_prim.GetAttribute("primvars:displayColor"))
        self.assertTrue(primvar)
        primvar.SetInterpolation(interpolation)

        await self._evaluate_graph(graph, controller)
        self._assert_no_compute_messages(graph, [og.WARNING, og.ERROR])

        target_prim = stage.GetPrimAtPath(target_prim_path)
        self.assertTrue(target_prim)
        attr = target_prim.GetAttribute("primvars:displayColor")
        self.assertTrue(attr)
        primvar = UsdGeom.Primvar(attr)
        self.assertTrue(primvar)
        self.assertEqual(primvar.GetInterpolation(), interpolation)

        # Finally, clear the metadata on the input prim and ensure that the metadata is cleared on the output prim
        attr = cube_prim.GetAttribute("primvars:displayColor")
        attr.ClearMetadata("interpolation")

        await self._evaluate_graph(graph, controller)
        # Just check for errors here; we do expect to get a warning that the metadata field no longer exists
        self._assert_no_compute_messages(graph, [og.ERROR])

        target_prim = stage.GetPrimAtPath(target_prim_path)
        self.assertTrue(target_prim)
        attr = target_prim.GetAttribute("primvars:displayColor")
        self.assertTrue(attr)
        self.assertFalse(attr.HasMetadata("interpolation"))

    # ----------------------------------------------------------------------
    # Helpers
    def _assert_no_compute_messages(self, graph, severity_list):
        graph_nodes = graph.get_nodes()
        for severity in severity_list:
            for node in graph_nodes:
                self.assertEqual(node.get_compute_messages(severity), [])

    async def _evaluate_graph(self, graph, controller):
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()

    def _get_node_messages(self, graph, severity):
        error_message = []
        graph_nodes = graph.get_nodes()
        for node in graph_nodes:
            error = node.get_compute_messages(severity)
            if error:
                error_message += error
        return error_message

    class CarbLogSuppressor:
        def __init__(self):
            self.logging = carb.logging.acquire_logging()
            self.log_was_enabled = self.logging.is_log_enabled()

        def __enter__(self):
            # turn off carb logging so the error does not trigger test failure
            # og logging still works and logs node error message.
            self.logging.set_log_enabled(False)
            return self

        def __exit__(self, *_):
            # restore carb logging
            self.logging.set_log_enabled(self.log_was_enabled)
