"""Tests that exercise node type forwarding using real node types as examples"""
import omni.graph.core as og
import omni.graph.core._unstable as ogu
import omni.graph.core.tests as ogts


# ==============================================================================================================
class TestNodeTypeForwarding(ogts.OmniGraphTestCase):
    async def test_simple_forward(self):
        """Test that a simple forwarding using a known hardcoded forwarding works properly"""
        old_type = "omni.graph.nodes.ComposeTuple2"
        new_type = "omni.graph.nodes.MakeVector2"
        new_extension = "omni.graph.nodes"

        interface = ogu.get_node_type_forwarding_interface()
        self.assertEqual(interface.find_forward(old_type, 1), (new_type, 1, new_extension))

        keys = og.Controller.Keys
        (_, (node,), _, _) = og.Controller.edit("/TestGraph", {keys.CREATE_NODES: ("ForwardedNode", old_type)})

        self.assertEqual(node.get_node_type().get_node_type(), new_type)

    # --------------------------------------------------------------------------------------------------------------
    async def test_runtime_forward(self):
        """Test that a forwarding that is added at runtime works properly"""
        old_type = "omni.graph.nodes.MakeVector2"
        new_type = "omni.graph.nodes.MakeVector3"
        new_extension = "omni.graph.nodes"

        interface = ogu.get_node_type_forwarding_interface()
        interface.define_forward(old_type, 1, new_type, 1, new_extension)
        self.assertEqual(interface.find_forward(old_type, 1), (new_type, 1, new_extension))

        keys = og.Controller.Keys
        (graph, (node,), _, _) = og.Controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: ("ForwardedNode", old_type),
            },
        )

        # Make sure the node type was properly forwarded
        self.assertEqual(node.get_node_type().get_node_type(), new_type)

        # Remove the forwarding and check that the same creation type now yields the original unforwarded node type
        interface.remove_forward(old_type, 1)
        (_, (new_node,), _, _) = og.Controller.edit(
            graph,
            {
                keys.CREATE_NODES: ("RegularNode", old_type),
            },
        )
        self.assertEqual(new_node.get_node_type().get_node_type(), old_type)
