From 1c6954950da4e2934b15c5f9010bcd17b83feae1 Mon Sep 17 00:00:00 2001
From: Arthur <worldnomad@gmail.com>
Date: Thu, 25 May 2023 12:02:45 -0700
Subject: [PATCH] 12552 test fixes

---
 netbox/utilities/testing/api.py   | 4 ++--
 netbox/utilities/testing/base.py  | 6 +++++-
 netbox/utilities/testing/views.py | 4 ++--
 netbox/utilities/tree_queries.py  | 8 ++++++--
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/netbox/utilities/testing/api.py b/netbox/utilities/testing/api.py
index 379e714fa..8135f731c 100644
--- a/netbox/utilities/testing/api.py
+++ b/netbox/utilities/testing/api.py
@@ -16,7 +16,7 @@ from users.models import ObjectPermission, Token
 from utilities.api import get_graphql_type_for_model
 from .base import ModelTestCase
 from .utils import disable_warnings
-
+from utilities.tree_queries import TreeQuerySet
 from ipam.graphql.types import IPAddressFamilyType
 
 
@@ -419,7 +419,7 @@ class APIViewTestCases:
 
             # Target the three most recently created objects to avoid triggering recursive deletions
             # (e.g. with django-tree-queries objects)
-            id_list = list(self._get_queryset().order_by('-id').values_list('id', flat=True)[:3])
+            id_list = list(self._get_queryset(with_tree_fields=False).order_by('-id').values_list('id', flat=True)[:3])
             self.assertEqual(len(id_list), 3, "Insufficient number of objects to test bulk deletion")
             data = [{"id": id} for id in id_list]
 
diff --git a/netbox/utilities/testing/base.py b/netbox/utilities/testing/base.py
index 04ceca1e2..761bd1ba4 100644
--- a/netbox/utilities/testing/base.py
+++ b/netbox/utilities/testing/base.py
@@ -12,6 +12,7 @@ from taggit.managers import TaggableManager
 
 from users.models import ObjectPermission
 from utilities.permissions import resolve_permission_ct
+from utilities.tree_queries import TreeQuerySet
 from utilities.utils import content_type_identifier
 from .utils import extract_form_failures
 
@@ -77,10 +78,13 @@ class ModelTestCase(TestCase):
     """
     model = None
 
-    def _get_queryset(self):
+    def _get_queryset(self, with_tree_fields=True):
         """
         Return a base queryset suitable for use in test methods.
         """
+        if not with_tree_fields and isinstance(self._get_queryset(), TreeQuerySet):
+            return self.model.objects.all().without_tree_fields()
+
         return self.model.objects.all()
 
     def prepare_instance(self, instance):
diff --git a/netbox/utilities/testing/views.py b/netbox/utilities/testing/views.py
index dc17548a2..a0888c759 100644
--- a/netbox/utilities/testing/views.py
+++ b/netbox/utilities/testing/views.py
@@ -507,7 +507,7 @@ class ViewTestCases:
             response = self.client.post(**request)
             self.assertHttpStatus(response, 302)
             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]:
+            for instance in self._get_queryset(with_tree_fields=False).order_by('-pk')[:self.bulk_create_count]:
                 self.assertInstanceEqual(instance, self.bulk_create_data, exclude=self.validation_excluded_fields)
 
         @override_settings(EXEMPT_VIEW_PERMISSIONS=[])
@@ -539,7 +539,7 @@ class ViewTestCases:
             response = self.client.post(**request)
             self.assertHttpStatus(response, 302)
             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]:
+            for instance in self._get_queryset(with_tree_fields=False).order_by('-pk')[:self.bulk_create_count]:
                 self.assertInstanceEqual(instance, self.bulk_create_data, exclude=self.validation_excluded_fields)
 
     class BulkImportObjectsViewTestCase(ModelViewTestCase):
diff --git a/netbox/utilities/tree_queries.py b/netbox/utilities/tree_queries.py
index e214838ed..59842e210 100644
--- a/netbox/utilities/tree_queries.py
+++ b/netbox/utilities/tree_queries.py
@@ -16,7 +16,11 @@ class TreeQuerySet(TreeQuerySet_, RestrictedQuerySet):
     Mate django-tree-queries TreeQuerySet with our RestrictedQuerySet for permissions enforcement.
     """
 
-    pass
+    def without_tree_fields(self):
+        """
+        Requests no tree fields on this queryset
+        """
+        return self.with_tree_fields(tree_fields=False)
 
 
 class TreeManager(Manager.from_queryset(TreeQuerySet), TreeManager_):
@@ -24,7 +28,7 @@ class TreeManager(Manager.from_queryset(TreeQuerySet), TreeManager_):
     Extend django-tree-queries TreeManager to incorporate RestrictedQuerySet().
     """
 
-    _with_tree_fields = False
+    _with_tree_fields = True
 
     def add_related_count(
         self,