"""
Tests related to omnigraph state/api when extensions get unloaded.
"""
from contextlib import suppress

import carb
import omni.graph.core as og
import omni.graph.core.tests as ogt
import omni.kit.app
import omni.kit.commands
import omni.kit.test
import omni.usd


# ======================================================================
class TestExtensionUnloads(ogt.OmniGraphTestCase):

    # -----------------------------------------------------------------------
    async def test_nodetype_when_extension_unloads(self):
        """Tests that nodes lose their node type when their extension unloads"""

        node_type = "omni.graph.examples.cpp.Simple"
        examples_cpp_extension = "omni.graph.examples.cpp"
        manager = omni.kit.app.get_app_interface().get_extension_manager()

        # If the extension is not found then the test cannot run
        examples_cpp_extension_id = None
        with suppress(Exception):
            examples_cpp_extension_id = manager.get_extension_id_by_module(examples_cpp_extension)
        if not examples_cpp_extension_id:
            carb.log_warn(
                f"test_nodetype_when_extension_unloads cannot run since {examples_cpp_extension} was not found"
            )
            return

        examples_cpp_enabled = manager.is_extension_enabled(examples_cpp_extension)
        try:
            manager.set_extension_enabled_immediate(examples_cpp_extension, True)

            # create a graph with a node from the extension
            keys = og.Controller.Keys
            controller1 = og.Controller()
            (_, nodes, _, _) = controller1.edit(
                "/World/Graph",
                {
                    keys.CREATE_NODES: [
                        ("Node", node_type),
                    ],
                },
            )
            self.assertTrue(nodes[0].get_node_type().is_valid())
            self.assertEquals(nodes[0].get_node_type().get_node_type(), node_type)

            # disable the extension, and validate the node type changes to a default type
            # the type is still valid, but is a default type
            manager.set_extension_enabled_immediate(examples_cpp_extension, False)
            self.assertTrue(nodes[0].get_node_type().is_valid())
            self.assertEquals(nodes[0].get_node_type().get_node_type(), None)

            # re-enable the extension and validate the node type is restored
            manager.set_extension_enabled_immediate(examples_cpp_extension, True)
            self.assertTrue(nodes[0].get_node_type().is_valid())
            self.assertEquals(nodes[0].get_node_type().get_node_type(), node_type)

        finally:
            manager.set_extension_enabled_immediate(examples_cpp_extension, examples_cpp_enabled)
