Update test methods to call unrestricted() on RestrictedQuerySets

netbox/utilities/testing/views.py

@@ -165,6 +165,14 @@ class ModelViewTestCase(TestCase):
         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.
+        """
+        if hasattr(self.model.objects, 'restrict'):
+            return self.model.objects.unrestricted()
+        return self.model.objects.all()
+
     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
@@ -208,12 +216,12 @@ class ViewTestCases:
         def test_get_object_anonymous(self):
             # Make the request as an unauthenticated user
             self.client.logout()
-            response = self.client.get(self.model.objects.first().get_absolute_url())
+            response = self.client.get(self._get_queryset().first().get_absolute_url())
             self.assertHttpStatus(response, 200)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
         def test_get_object_without_permission(self):
-            instance = self.model.objects.first()
+            instance = self._get_queryset().first()
 
             # Try GET without permission
             with disable_warnings('django.request'):
@@ -221,7 +229,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
         def test_get_object_with_permission(self):
-            instance = self.model.objects.first()
+            instance = self._get_queryset().first()
 
             # Add model-level permission
             obj_perm = ObjectPermission(
@@ -236,7 +244,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
         def test_get_object_with_constrained_permission(self):
-            instance1, instance2 = self.model.objects.all()[:2]
+            instance1, instance2 = self._get_queryset().all()[:2]
 
             # Add object-level permission
             obj_perm = ObjectPermission(
@@ -259,7 +267,7 @@ class ViewTestCases:
         """
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_get_object_changelog(self):
-            url = self._get_url('changelog', self.model.objects.first())
+            url = self._get_url('changelog', self._get_queryset().first())
             response = self.client.get(url)
             self.assertHttpStatus(response, 200)
 
@@ -288,7 +296,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_create_object_with_permission(self):
-            initial_count = self.model.objects.count()
+            initial_count = self._get_queryset().count()
 
             # Assign unconstrained permission
             obj_perm = ObjectPermission(
@@ -307,12 +315,12 @@ class ViewTestCases:
                 'data': post_data(self.form_data),
             }
             self.assertHttpStatus(self.client.post(**request), 302)
-            self.assertEqual(initial_count + 1, self.model.objects.count())
-            self.assertInstanceEqual(self.model.objects.order_by('pk').last(), self.form_data)
+            self.assertEqual(initial_count + 1, self._get_queryset().count())
+            self.assertInstanceEqual(self._get_queryset().order_by('pk').last(), self.form_data)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_create_object_with_constrained_permission(self):
-            initial_count = self.model.objects.count()
+            initial_count = self._get_queryset().count()
 
             # Assign constrained permission
             obj_perm = ObjectPermission(
@@ -332,7 +340,7 @@ class ViewTestCases:
                 'data': post_data(self.form_data),
             }
             self.assertHttpStatus(self.client.post(**request), 200)
-            self.assertEqual(initial_count, self.model.objects.count())  # Check that no object was created
+            self.assertEqual(initial_count, self._get_queryset().count())  # Check that no object was created
 
             # Update the ObjectPermission to allow creation
             obj_perm.constraints = {'pk__gt': 0}
@@ -344,8 +352,8 @@ class ViewTestCases:
                 'data': post_data(self.form_data),
             }
             self.assertHttpStatus(self.client.post(**request), 302)
-            self.assertEqual(initial_count + 1, self.model.objects.count())
-            self.assertInstanceEqual(self.model.objects.order_by('pk').last(), self.form_data)
+            self.assertEqual(initial_count + 1, self._get_queryset().count())
+            self.assertInstanceEqual(self._get_queryset().order_by('pk').last(), self.form_data)
 
     class EditObjectViewTestCase(ModelViewTestCase):
         """
@@ -356,7 +364,7 @@ class ViewTestCases:
         form_data = {}
 
         def test_edit_object_without_permission(self):
-            instance = self.model.objects.first()
+            instance = self._get_queryset().first()
 
             # Try GET without permission
             with disable_warnings('django.request'):
@@ -372,7 +380,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_edit_object_with_permission(self):
-            instance = self.model.objects.first()
+            instance = self._get_queryset().first()
 
             # Assign model-level permission
             obj_perm = ObjectPermission(
@@ -391,11 +399,11 @@ class ViewTestCases:
                 'data': post_data(self.form_data),
             }
             self.assertHttpStatus(self.client.post(**request), 302)
-            self.assertInstanceEqual(self.model.objects.get(pk=instance.pk), self.form_data)
+            self.assertInstanceEqual(self._get_queryset().get(pk=instance.pk), self.form_data)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_edit_object_with_constrained_permission(self):
-            instance1, instance2 = self.model.objects.all()[:2]
+            instance1, instance2 = self._get_queryset().all()[:2]
 
             # Assign constrained permission
             obj_perm = ObjectPermission(
@@ -418,7 +426,7 @@ class ViewTestCases:
                 'data': post_data(self.form_data),
             }
             self.assertHttpStatus(self.client.post(**request), 302)
-            self.assertInstanceEqual(self.model.objects.get(pk=instance1.pk), self.form_data)
+            self.assertInstanceEqual(self._get_queryset().get(pk=instance1.pk), self.form_data)
 
             # Try to edit a non-permitted object
             request = {
@@ -432,7 +440,7 @@ class ViewTestCases:
         Delete a single instance.
         """
         def test_delete_object_without_permission(self):
-            instance = self.model.objects.first()
+            instance = self._get_queryset().first()
 
             # Try GET without permission
             with disable_warnings('django.request'):
@@ -448,7 +456,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_delete_object_with_permission(self):
-            instance = self.model.objects.first()
+            instance = self._get_queryset().first()
 
             # Assign model-level permission
             obj_perm = ObjectPermission(
@@ -468,11 +476,11 @@ class ViewTestCases:
             }
             self.assertHttpStatus(self.client.post(**request), 302)
             with self.assertRaises(ObjectDoesNotExist):
-                self.model.objects.get(pk=instance.pk)
+                self._get_queryset().get(pk=instance.pk)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_delete_object_with_constrained_permission(self):
-            instance1, instance2 = self.model.objects.all()[:2]
+            instance1, instance2 = self._get_queryset().all()[:2]
 
             # Assign object-level permission
             obj_perm = ObjectPermission(
@@ -496,7 +504,7 @@ class ViewTestCases:
             }
             self.assertHttpStatus(self.client.post(**request), 302)
             with self.assertRaises(ObjectDoesNotExist):
-                self.model.objects.get(pk=instance1.pk)
+                self._get_queryset().get(pk=instance1.pk)
 
             # Try to delete a non-permitted object
             request = {
@@ -504,7 +512,7 @@ class ViewTestCases:
                 'data': post_data({'confirm': True}),
             }
             self.assertHttpStatus(self.client.post(**request), 404)
-            self.assertTrue(self.model.objects.filter(pk=instance2.pk).exists())
+            self.assertTrue(self._get_queryset().filter(pk=instance2.pk).exists())
 
     class ListObjectsViewTestCase(ModelViewTestCase):
         """
@@ -546,7 +554,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
         def test_list_objects_with_constrained_permission(self):
-            instance1, instance2 = self.model.objects.all()[:2]
+            instance1, instance2 = self._get_queryset().all()[:2]
 
             # Add object-level permission
             obj_perm = ObjectPermission(
@@ -591,7 +599,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
         def test_bulk_create_objects_with_permission(self):
-            initial_count = self.model.objects.count()
+            initial_count = self._get_queryset().count()
             request = {
                 'path': self._get_url('add'),
                 'data': post_data(self.bulk_create_data),
@@ -608,13 +616,13 @@ class ViewTestCases:
             # Bulk create objects
             response = self.client.post(**request)
             self.assertHttpStatus(response, 302)
-            self.assertEqual(initial_count + self.bulk_create_count, self.model.objects.count())
-            for instance in self.model.objects.order_by('-pk')[:self.bulk_create_count]:
+            self.assertEqual(initial_count + self.bulk_create_count, self._get_queryset().count())
+            for instance in self._get_queryset().order_by('-pk')[:self.bulk_create_count]:
                 self.assertInstanceEqual(instance, self.bulk_create_data)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
         def test_bulk_create_objects_with_constrained_permission(self):
-            initial_count = self.model.objects.count()
+            initial_count = self._get_queryset().count()
             request = {
                 'path': self._get_url('add'),
                 'data': post_data(self.bulk_create_data),
@@ -631,7 +639,7 @@ class ViewTestCases:
 
             # Attempt to make the request with unmet constraints
             self.assertHttpStatus(self.client.post(**request), 200)
-            self.assertEqual(self.model.objects.count(), initial_count)
+            self.assertEqual(self._get_queryset().count(), initial_count)
 
             # Update the ObjectPermission to allow creation
             obj_perm.constraints = {'pk__gt': 0}  # Dummy constraint to allow all
@@ -639,8 +647,8 @@ class ViewTestCases:
 
             response = self.client.post(**request)
             self.assertHttpStatus(response, 302)
-            self.assertEqual(initial_count + self.bulk_create_count, self.model.objects.count())
-            for instance in self.model.objects.order_by('-pk')[:self.bulk_create_count]:
+            self.assertEqual(initial_count + self.bulk_create_count, self._get_queryset().count())
+            for instance in self._get_queryset().order_by('-pk')[:self.bulk_create_count]:
                 self.assertInstanceEqual(instance, self.bulk_create_data)
 
     class BulkImportObjectsViewTestCase(ModelViewTestCase):
@@ -670,7 +678,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_bulk_import_objects_with_permission(self):
-            initial_count = self.model.objects.count()
+            initial_count = self._get_queryset().count()
             data = {
                 'csv': self._get_csv_data(),
             }
@@ -688,11 +696,11 @@ class ViewTestCases:
 
             # Test POST with permission
             self.assertHttpStatus(self.client.post(self._get_url('import'), data), 200)
-            self.assertEqual(self.model.objects.count(), initial_count + len(self.csv_data) - 1)
+            self.assertEqual(self._get_queryset().count(), initial_count + len(self.csv_data) - 1)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_bulk_import_objects_with_constrained_permission(self):
-            initial_count = self.model.objects.count()
+            initial_count = self._get_queryset().count()
             data = {
                 'csv': self._get_csv_data(),
             }
@@ -708,7 +716,7 @@ class ViewTestCases:
 
             # Attempt to import non-permitted objects
             self.assertHttpStatus(self.client.post(self._get_url('import'), data), 200)
-            self.assertEqual(self.model.objects.count(), initial_count)
+            self.assertEqual(self._get_queryset().count(), initial_count)
 
             # Update permission constraints
             obj_perm.constraints = {'pk__gt': 0}  # Dummy permission to allow all
@@ -716,7 +724,7 @@ class ViewTestCases:
 
             # Import permitted objects
             self.assertHttpStatus(self.client.post(self._get_url('import'), data), 200)
-            self.assertEqual(self.model.objects.count(), initial_count + len(self.csv_data) - 1)
+            self.assertEqual(self._get_queryset().count(), initial_count + len(self.csv_data) - 1)
 
     class BulkEditObjectsViewTestCase(ModelViewTestCase):
         """
@@ -728,7 +736,7 @@ class ViewTestCases:
         bulk_edit_data = {}
 
         def test_bulk_edit_objects_without_permission(self):
-            pk_list = self.model.objects.values_list('pk', flat=True)[:3]
+            pk_list = self._get_queryset().values_list('pk', flat=True)[:3]
             data = {
                 'pk': pk_list,
                 '_apply': True,  # Form button
@@ -744,7 +752,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_bulk_edit_objects_with_permission(self):
-            pk_list = self.model.objects.values_list('pk', flat=True)[:3]
+            pk_list = self._get_queryset().values_list('pk', flat=True)[:3]
             data = {
                 'pk': pk_list,
                 '_apply': True,  # Form button
@@ -763,13 +771,12 @@ class ViewTestCases:
 
             # Try POST with model-level permission
             self.assertHttpStatus(self.client.post(self._get_url('bulk_edit'), data), 302)
-            for i, instance in enumerate(self.model.objects.filter(pk__in=pk_list)):
+            for i, instance in enumerate(self._get_queryset().filter(pk__in=pk_list)):
                 self.assertInstanceEqual(instance, self.bulk_edit_data)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=['*'])
         def test_bulk_edit_objects_with_constrained_permission(self):
-            initial_instances = self.model.objects.all()[:3]
-            pk_list = list(self.model.objects.values_list('pk', flat=True)[:3])
+            pk_list = list(self._get_queryset().values_list('pk', flat=True)[:3])
             data = {
                 'pk': pk_list,
                 '_apply': True,  # Form button
@@ -781,7 +788,7 @@ class ViewTestCases:
             # Dynamically determine a constraint that will *not* be matched by the updated objects.
             attr_name = list(self.bulk_edit_data.keys())[0]
             field = self.model._meta.get_field(attr_name)
-            value = field.value_from_object(self.model.objects.first())
+            value = field.value_from_object(self._get_queryset().first())
 
             # Assign constrained permission
             obj_perm = ObjectPermission(
@@ -802,7 +809,7 @@ class ViewTestCases:
 
             # Bulk edit permitted objects
             self.assertHttpStatus(self.client.post(self._get_url('bulk_edit'), data), 302)
-            for i, instance in enumerate(self.model.objects.filter(pk__in=pk_list)):
+            for i, instance in enumerate(self._get_queryset().filter(pk__in=pk_list)):
                 self.assertInstanceEqual(instance, self.bulk_edit_data)
 
     class BulkDeleteObjectsViewTestCase(ModelViewTestCase):
@@ -811,7 +818,7 @@ class ViewTestCases:
         """
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
         def test_bulk_delete_objects_without_permission(self):
-            pk_list = self.model.objects.values_list('pk', flat=True)[:3]
+            pk_list = self._get_queryset().values_list('pk', flat=True)[:3]
             data = {
                 'pk': pk_list,
                 'confirm': True,
@@ -828,7 +835,7 @@ class ViewTestCases:
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
         def test_bulk_delete_objects_with_permission(self):
-            pk_list = self.model.objects.values_list('pk', flat=True)
+            pk_list = self._get_queryset().values_list('pk', flat=True)
             data = {
                 'pk': pk_list,
                 'confirm': True,
@@ -845,12 +852,12 @@ class ViewTestCases:
 
             # Try POST with model-level permission
             self.assertHttpStatus(self.client.post(self._get_url('bulk_delete'), data), 302)
-            self.assertEqual(self.model.objects.count(), 0)
+            self.assertEqual(self._get_queryset().count(), 0)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
         def test_bulk_delete_objects_with_constrained_permission(self):
-            initial_count = self.model.objects.count()
-            pk_list = self.model.objects.values_list('pk', flat=True)
+            initial_count = self._get_queryset().count()
+            pk_list = self._get_queryset().values_list('pk', flat=True)
             data = {
                 'pk': pk_list,
                 'confirm': True,
@@ -868,7 +875,7 @@ class ViewTestCases:
 
             # Attempt to bulk delete non-permitted objects
             self.assertHttpStatus(self.client.post(self._get_url('bulk_delete'), data), 302)
-            self.assertEqual(self.model.objects.count(), initial_count)
+            self.assertEqual(self._get_queryset().count(), initial_count)
 
             # Update permission constraints
             obj_perm.constraints = {'pk__gt': 0}  # Dummy permission to allow all
@@ -876,7 +883,7 @@ class ViewTestCases:
 
             # Bulk delete permitted objects
             self.assertHttpStatus(self.client.post(self._get_url('bulk_delete'), data), 302)
-            self.assertEqual(self.model.objects.count(), 0)
+            self.assertEqual(self._get_queryset().count(), 0)
 
     class PrimaryObjectViewTestCase(
         GetObjectViewTestCase,