Merge pull request #134 from netboxlabs/fix-testing

Fix get_models for testing

Author: bctiemann

netbox_custom_objects/__init__.py

@@ -1,6 +1,9 @@
 import sys
+import warnings
 
 from django.core.exceptions import AppRegistryNotReady
+from django.db import transaction
+from django.db.utils import DatabaseError, OperationalError, ProgrammingError
 from netbox.plugins import PluginConfig
 
 
@@ -9,12 +12,32 @@ def is_running_migration():
     Check if the code is currently running during a Django migration.
     """
     # Check if 'makemigrations' or 'migrate' command is in sys.argv
-    if any(cmd in sys.argv for cmd in ['makemigrations', 'migrate']):
+    if any(cmd in sys.argv for cmd in ["makemigrations", "migrate"]):
         return True
 
     return False
 
 
+def check_custom_object_type_table_exists():
+    """
+    Check if the CustomObjectType table exists in the database.
+    Returns True if the table exists, False otherwise.
+    """
+    from .models import CustomObjectType
+
+    try:
+        # Try to query the model - if the table doesn't exist, this will raise an exception
+        # this check and the transaction.atomic() is only required when running tests as the
+        # migration check doesn't work correctly in the test environment
+        with transaction.atomic():
+            # Force immediate execution by using first()
+            CustomObjectType.objects.first()
+        return True
+    except (OperationalError, ProgrammingError, DatabaseError):
+        # Catch database-specific errors (table doesn't exist, permission issues, etc.)
+        return False
+
+
 # Plugin Configuration
 class CustomObjectsPluginConfig(PluginConfig):
     name = "netbox_custom_objects"
@@ -59,37 +82,35 @@ def get_model(self, model_name, require_ready=True):
 
         return obj.get_model()
 
-    # def get_models(self, include_auto_created=False, include_swapped=False):
-    #     """Return all models for this plugin, including custom object type models."""
-    #     # Get the regular Django models first
-    #     for model in super().get_models(include_auto_created, include_swapped):
-    #         yield model
-    #
-    #     # Skip custom object type model loading if running during migration
-    #     if is_running_migration():
-    #         return
-    #
-    #     # Suppress warnings about database calls during model loading
-    #     with warnings.catch_warnings():
-    #         warnings.filterwarnings(
-    #             "ignore", category=RuntimeWarning, message=".*database.*"
-    #         )
-    #         warnings.filterwarnings(
-    #             "ignore", category=UserWarning, message=".*database.*"
-    #         )
-    #
-    #         # Add custom object type models
-    #         from .models import CustomObjectType
-    #
-    #         # Only load models that are already cached to avoid creating all models at startup
-    #         # This prevents the "two TaggableManagers with same through model" error
-    #         custom_object_types = CustomObjectType.objects.all()
-    #         for custom_type in custom_object_types:
-    #             # Only yield already cached models during discovery
-    #             if CustomObjectType.is_model_cached(custom_type.id):
-    #                 model = CustomObjectType.get_cached_model(custom_type.id)
-    #                 if model:
-    #                     yield model
+    def get_models(self, include_auto_created=False, include_swapped=False):
+        """Return all models for this plugin, including custom object type models."""
+        # Get the regular Django models first
+        for model in super().get_models(include_auto_created, include_swapped):
+            yield model
+
+        # Suppress warnings about database calls during model loading
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore", category=RuntimeWarning, message=".*database.*"
+            )
+            warnings.filterwarnings(
+                "ignore", category=UserWarning, message=".*database.*"
+            )
+
+            # Skip custom object type model loading if running during migration
+            if is_running_migration() or not check_custom_object_type_table_exists():
+                return
+
+            # Add custom object type models
+            from .models import CustomObjectType
+
+            custom_object_types = CustomObjectType.objects.all()
+            for custom_type in custom_object_types:
+                # Only yield already cached models during discovery
+                if CustomObjectType.is_model_cached(custom_type.id):
+                    model = CustomObjectType.get_cached_model(custom_type.id)
+                    if model:
+                        yield model
 
 
 config = CustomObjectsPluginConfig