Refactor APITestCase to provide dynamic queryset specification

netbox/utilities/testing/api.py

@@ -7,7 +7,7 @@ from rest_framework.test import APIClient
 
 from users.models import ObjectPermission, Token
 from .utils import disable_warnings
-from .views import TestCase
+from .views import ModelTestCase
 
 
 __all__ = (
@@ -20,9 +20,8 @@ __all__ = (
 # REST API Tests
 #
 
-class APITestCase(TestCase):
+class APITestCase(ModelTestCase):
     client_class = APIClient
-    model = None
 
     def setUp(self):
         """
@@ -52,7 +51,7 @@ class APIViewTestCases:
             """
             GET a single object as an unauthenticated user.
             """
-            url = self._get_detail_url(self.model.objects.unrestricted().first())
+            url = self._get_detail_url(self._get_queryset().first())
             response = self.client.get(url, **self.header)
             self.assertHttpStatus(response, status.HTTP_200_OK)
 
@@ -61,7 +60,7 @@ class APIViewTestCases:
             """
             GET a single object as an authenticated user without the required permission.
             """
-            url = self._get_detail_url(self.model.objects.unrestricted().first())
+            url = self._get_detail_url(self._get_queryset().first())
 
             # Try GET without permission
             with disable_warnings('django.request'):
@@ -72,9 +71,9 @@ class APIViewTestCases:
             """
             GET a single object as an authenticated user with permission to view the object.
             """
-            self.assertGreaterEqual(self.model.objects.unrestricted().count(), 2,
+            self.assertGreaterEqual(self._get_queryset().count(), 2,
                                     f"Test requires the creation of at least two {self.model} instances")
-            instance1, instance2 = self.model.objects.unrestricted()[:2]
+            instance1, instance2 = self._get_queryset()[:2]
 
             # Add object-level permission
             obj_perm = ObjectPermission(
@@ -104,7 +103,7 @@ class APIViewTestCases:
             url = self._get_list_url()
             response = self.client.get(url, **self.header)
 
-            self.assertEqual(len(response.data['results']), self.model.objects.unrestricted().count())
+            self.assertEqual(len(response.data['results']), self._get_queryset().count())
             self.assertHttpStatus(response, status.HTTP_200_OK)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
@@ -115,7 +114,7 @@ class APIViewTestCases:
             url = f'{self._get_list_url()}?brief=1'
             response = self.client.get(url, **self.header)
 
-            self.assertEqual(len(response.data['results']), self.model.objects.unrestricted().count())
+            self.assertEqual(len(response.data['results']), self._get_queryset().count())
             self.assertEqual(sorted(response.data['results'][0]), self.brief_fields)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
@@ -134,9 +133,9 @@ class APIViewTestCases:
             """
             GET a list of objects as an authenticated user with permission to view the objects.
             """
-            self.assertGreaterEqual(self.model.objects.unrestricted().count(), 3,
+            self.assertGreaterEqual(self._get_queryset().count(), 3,
                                     f"Test requires the creation of at least three {self.model} instances")
-            instance1, instance2 = self.model.objects.unrestricted()[:2]
+            instance1, instance2 = self._get_queryset()[:2]
 
             # Add object-level permission
             obj_perm = ObjectPermission(
@@ -178,12 +177,12 @@ class APIViewTestCases:
             obj_perm.users.add(self.user)
             obj_perm.object_types.add(ContentType.objects.get_for_model(self.model))
 
-            initial_count = self.model.objects.unrestricted().count()
+            initial_count = self._get_queryset().count()
             response = self.client.post(self._get_list_url(), self.create_data[0], format='json', **self.header)
             self.assertHttpStatus(response, status.HTTP_201_CREATED)
-            self.assertEqual(self.model.objects.unrestricted().count(), initial_count + 1)
+            self.assertEqual(self._get_queryset().count(), initial_count + 1)
             self.assertInstanceEqual(
-                self.model.objects.unrestricted().get(pk=response.data['id']),
+                self._get_queryset().get(pk=response.data['id']),
                 self.create_data[0],
                 api=True
             )
@@ -200,14 +199,14 @@ class APIViewTestCases:
             obj_perm.users.add(self.user)
             obj_perm.object_types.add(ContentType.objects.get_for_model(self.model))
 
-            initial_count = self.model.objects.unrestricted().count()
+            initial_count = self._get_queryset().count()
             response = self.client.post(self._get_list_url(), self.create_data, format='json', **self.header)
             self.assertHttpStatus(response, status.HTTP_201_CREATED)
             self.assertEqual(len(response.data), len(self.create_data))
-            self.assertEqual(self.model.objects.unrestricted().count(), initial_count + len(self.create_data))
+            self.assertEqual(self._get_queryset().count(), initial_count + len(self.create_data))
             for i, obj in enumerate(response.data):
                 self.assertInstanceEqual(
-                    self.model.objects.unrestricted().get(pk=obj['id']),
+                    self._get_queryset().get(pk=obj['id']),
                     self.create_data[i],
                     api=True
                 )
@@ -219,7 +218,7 @@ class APIViewTestCases:
             """
             PATCH a single object without permission.
             """
-            url = self._get_detail_url(self.model.objects.unrestricted().first())
+            url = self._get_detail_url(self._get_queryset().first())
             update_data = self.update_data or getattr(self, 'create_data')[0]
 
             # Try PATCH without permission
@@ -231,7 +230,7 @@ class APIViewTestCases:
             """
             PATCH a single object identified by its numeric ID.
             """
-            instance = self.model.objects.unrestricted().first()
+            instance = self._get_queryset().first()
             url = self._get_detail_url(instance)
             update_data = self.update_data or getattr(self, 'create_data')[0]
 
@@ -254,7 +253,7 @@ class APIViewTestCases:
             """
             DELETE a single object without permission.
             """
-            url = self._get_detail_url(self.model.objects.unrestricted().first())
+            url = self._get_detail_url(self._get_queryset().first())
 
             # Try DELETE without permission
             with disable_warnings('django.request'):
@@ -265,7 +264,7 @@ class APIViewTestCases:
             """
             DELETE a single object identified by its numeric ID.
             """
-            instance = self.model.objects.unrestricted().first()
+            instance = self._get_queryset().first()
             url = self._get_detail_url(instance)
 
             # Add object-level permission
@@ -278,7 +277,7 @@ class APIViewTestCases:
 
             response = self.client.delete(url, **self.header)
             self.assertHttpStatus(response, status.HTTP_204_NO_CONTENT)
-            self.assertFalse(self.model.objects.unrestricted().filter(pk=instance.pk).exists())
+            self.assertFalse(self._get_queryset().filter(pk=instance.pk).exists())
 
     class APIViewTestCase(
         GetObjectViewTestCase,

netbox/utilities/testing/views.py

@@ -18,6 +18,7 @@ from .utils import disable_warnings, post_data
 
 __all__ = (
     'TestCase',
+    'ModelTestCase',
     'ModelViewTestCase',
     'ViewTestCases',
 )
@@ -149,22 +150,12 @@ class TestCase(_TestCase):
         return tags
 
 
-#
-# UI Tests
-#
-
-class ModelViewTestCase(TestCase):
+class ModelTestCase(TestCase):
     """
-    Base TestCase for model views. Subclass to test individual views.
+    Parent class for TestCases which deal with models.
     """
     model = None
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        if self.model is None:
-            raise Exception("Test case requires model to be defined")
-
     def _get_queryset(self):
         """
         Return a base queryset suitable for use in test methods. Call unrestricted() if RestrictedQuerySet is in use.
@@ -173,6 +164,16 @@ class ModelViewTestCase(TestCase):
             return self.model.objects.unrestricted()
         return self.model.objects.all()
 
+
+#
+# UI Tests
+#
+
+class ModelViewTestCase(ModelTestCase):
+    """
+    Base TestCase for model views. Subclass to test individual views.
+    """
+
     def _get_base_url(self):
         """
         Return the base format for a URL for the test's model. Override this to test for a model which belongs