John Anderson 5 лет назад
Родитель
Сommit
9e84e3b83b
3 измененных файлов с 295 добавлено и 10 удалено
  1. 6 7
      netbox/extras/querysets.py
  2. 278 2
      netbox/extras/tests/test_models.py
  3. 11 1
      netbox/utilities/query_functions.py

+ 6 - 7
netbox/extras/querysets.py

@@ -1,9 +1,9 @@
 from collections import OrderedDict
 
 from django.contrib.postgres.aggregates import JSONBAgg
-from django.db.models import OuterRef, Subquery, Q, QuerySet
+from django.db.models import OuterRef, Subquery, Q, QuerySet, OrderBy, F, ExpressionWrapper
 
-from utilities.query_functions import EmptyGroupByJSONBAgg
+from utilities.query_functions import EmptyGroupByJSONBAgg, OrderableJSONBAgg
 from utilities.querysets import RestrictedQuerySet
 
 
@@ -64,7 +64,9 @@ class ConfigContextQuerySet(RestrictedQuerySet):
         ).order_by('weight', 'name')
 
         if aggregate_data:
-            queryset = queryset.aggregate(config_context_data=JSONBAgg('data'))['config_context_data']
+            return queryset.aggregate(
+                config_context_data=OrderableJSONBAgg('data', ordering=['weight', 'name'])
+            )['config_context_data']
 
         return queryset
 
@@ -90,11 +92,8 @@ class ConfigContextModelQuerySet(RestrictedQuerySet):
             config_context_data=Subquery(
                 ConfigContext.objects.filter(
                     self._get_config_context_filters()
-                ).order_by(
-                    'weight',
-                    'name'
                 ).annotate(
-                    _data=EmptyGroupByJSONBAgg('data')
+                    _data=EmptyGroupByJSONBAgg('data', ordering=['weight', 'name'])
                 ).values("_data")
             )
         )

+ 278 - 2
netbox/extras/tests/test_models.py

@@ -1,9 +1,11 @@
 from django.contrib.contenttypes.models import ContentType
 from django.test import TestCase
 
-from dcim.models import Site
+from dcim.models import Device, DeviceRole, DeviceType, Manufacturer, Platform, Site, Region
 from extras.choices import TemplateLanguageChoices
-from extras.models import Graph, Tag
+from extras.models import ConfigContext, Graph, Tag
+from tenancy.models import Tenant, TenantGroup
+from virtualization.models import Cluster, ClusterGroup, ClusterType, VirtualMachine
 
 
 class GraphTest(TestCase):
@@ -53,3 +55,277 @@ class TagTest(TestCase):
         tag.save()
 
         self.assertEqual(tag.slug, 'testing-unicode-台灣')
+
+
+class ConfigContextTest(TestCase):
+    """
+    These test cases deal with the weighting, ordering, and deep merge logic of config context data.
+
+    It also ensures the various config context querysets are consistent.
+    """
+
+    def setUp(self):
+
+        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        self.devicetype = DeviceType.objects.create(manufacturer=manufacturer, model='Device Type 1', slug='device-type-1')
+        self.devicerole = DeviceRole.objects.create(name='Device Role 1', slug='device-role-1')
+        self.region = Region.objects.create(name="Region")
+        self.site = Site.objects.create(name='Site-1', slug='site-1', region=self.region)
+        self.platform = Platform.objects.create(name="Platform")
+        self.tenantgroup = TenantGroup.objects.create(name="Tenant Group")
+        self.tenant = Tenant.objects.create(name="Tenant", group=self.tenantgroup)
+        self.tag = Tag.objects.create(name="Tag", slug="tag")
+
+        self.device = Device.objects.create(
+            name='Device 1',
+            device_type=self.devicetype,
+            device_role=self.devicerole,
+            site=self.site
+        )
+
+    def test_higher_weight_wins(self):
+
+        context1 = ConfigContext(
+            name="context 1",
+            weight=101,
+            data={
+                "a": 123,
+                "b": 456,
+                "c": 777
+            }
+        )
+        context2 = ConfigContext(
+            name="context 2",
+            weight=100,
+            data={
+                "a": 123,
+                "b": 456,
+                "c": 789
+            }
+        )
+        ConfigContext.objects.bulk_create([context1, context2])
+
+        expected_data = {
+            "a": 123,
+            "b": 456,
+            "c": 777
+        }
+        self.assertEqual(self.device.get_config_context(), expected_data)
+
+    def test_name_ordering_after_weight(self):
+
+        context1 = ConfigContext(
+            name="context 1",
+            weight=100,
+            data={
+                "a": 123,
+                "b": 456,
+                "c": 777
+            }
+        )
+        context2 = ConfigContext(
+            name="context 2",
+            weight=100,
+            data={
+                "a": 123,
+                "b": 456,
+                "c": 789
+            }
+        )
+        ConfigContext.objects.bulk_create([context1, context2])
+
+        expected_data = {
+            "a": 123,
+            "b": 456,
+            "c": 789
+        }
+        self.assertEqual(self.device.get_config_context(), expected_data)
+
+    def test_annotation_same_as_get_for_object(self):
+        """
+        This test incorperates features from all of the above tests cases to ensure
+        the annotate_config_context_data() and get_for_object() queryset methods are the same.
+        """
+        context1 = ConfigContext(
+            name="context 1",
+            weight=101,
+            data={
+                "a": 123,
+                "b": 456,
+                "c": 777
+            }
+        )
+        context2 = ConfigContext(
+            name="context 2",
+            weight=100,
+            data={
+                "a": 123,
+                "b": 456,
+                "c": 789
+            }
+        )
+        context3 = ConfigContext(
+            name="context 3",
+            weight=99,
+            data={
+                "d": 1
+            }
+        )
+        context4 = ConfigContext(
+            name="context 4",
+            weight=99,
+            data={
+                "d": 2
+            }
+        )
+        ConfigContext.objects.bulk_create([context1, context2, context3, context4])
+
+        annotated_queryset = Device.objects.filter(name=self.device.name).annotate_config_context_data()
+        self.assertEqual(self.device.get_config_context(), annotated_queryset[0].get_config_context())
+
+    def test_annotation_same_as_get_for_object_device_relations(self):
+
+        site_context = ConfigContext.objects.create(
+            name="site",
+            weight=100,
+            data={
+                "site": 1
+            }
+        )
+        site_context.sites.add(self.site)
+        region_context = ConfigContext.objects.create(
+            name="region",
+            weight=100,
+            data={
+                "region": 1
+            }
+        )
+        region_context.regions.add(self.region)
+        platform_context = ConfigContext.objects.create(
+            name="platform",
+            weight=100,
+            data={
+                "platform": 1
+            }
+        )
+        platform_context.platforms.add(self.platform)
+        tenant_group_context = ConfigContext.objects.create(
+            name="tenant group",
+            weight=100,
+            data={
+                "tenant_group": 1
+            }
+        )
+        tenant_group_context.tenant_groups.add(self.tenantgroup)
+        tenant_context = ConfigContext.objects.create(
+            name="tenant",
+            weight=100,
+            data={
+                "tenant": 1
+            }
+        )
+        tenant_context.tenants.add(self.tenant)
+        tag_context = ConfigContext.objects.create(
+            name="tag",
+            weight=100,
+            data={
+                "tag": 1
+            }
+        )
+        tag_context.tags.add(self.tag)
+
+        device = Device.objects.create(
+            name="Device 2",
+            site=self.site,
+            tenant=self.tenant,
+            platform=self.platform,
+            device_role=self.devicerole,
+            device_type=self.devicetype
+        )
+        device.tags.add(self.tag)
+
+        annotated_queryset = Device.objects.filter(name=device.name).annotate_config_context_data()
+        self.assertEqual(device.get_config_context(), annotated_queryset[0].get_config_context())
+
+    def test_annotation_same_as_get_for_object_virtualmachine_relations(self):
+
+        site_context = ConfigContext.objects.create(
+            name="site",
+            weight=100,
+            data={
+                "site": 1
+            }
+        )
+        site_context.sites.add(self.site)
+        region_context = ConfigContext.objects.create(
+            name="region",
+            weight=100,
+            data={
+                "region": 1
+            }
+        )
+        region_context.regions.add(self.region)
+        platform_context = ConfigContext.objects.create(
+            name="platform",
+            weight=100,
+            data={
+                "platform": 1
+            }
+        )
+        platform_context.platforms.add(self.platform)
+        tenant_group_context = ConfigContext.objects.create(
+            name="tenant group",
+            weight=100,
+            data={
+                "tenant_group": 1
+            }
+        )
+        tenant_group_context.tenant_groups.add(self.tenantgroup)
+        tenant_context = ConfigContext.objects.create(
+            name="tenant",
+            weight=100,
+            data={
+                "tenant": 1
+            }
+        )
+        tenant_context.tenants.add(self.tenant)
+        tag_context = ConfigContext.objects.create(
+            name="tag",
+            weight=100,
+            data={
+                "tag": 1
+            }
+        )
+        tag_context.tags.add(self.tag)
+        cluster_group = ClusterGroup.objects.create(name="Cluster Group")
+        cluster_group_context = ConfigContext.objects.create(
+            name="cluster group",
+            weight=100,
+            data={
+                "cluster_group": 1
+            }
+        )
+        cluster_group_context.cluster_groups.add(cluster_group)
+        cluster_type = ClusterType.objects.create(name="Cluster Type 1")
+        cluster = Cluster.objects.create(name="Cluster", group=cluster_group, type=cluster_type)
+        cluster_context = ConfigContext.objects.create(
+            name="cluster",
+            weight=100,
+            data={
+                "cluster": 1
+            }
+        )
+        cluster_context.clusters.add(cluster)
+
+
+        virtual_machine = VirtualMachine.objects.create(
+            name="VM 1",
+            cluster=cluster,
+            tenant=self.tenant,
+            platform=self.platform,
+            role=self.devicerole
+        )
+        virtual_machine.tags.add(self.tag)
+
+        annotated_queryset = VirtualMachine.objects.filter(name=virtual_machine.name).annotate_config_context_data()
+        self.assertEqual(virtual_machine.get_config_context(), annotated_queryset[0].get_config_context())

+ 11 - 1
netbox/utilities/query_functions.py

@@ -1,4 +1,5 @@
 from django.contrib.postgres.aggregates import JSONBAgg
+from django.contrib.postgres.aggregates.mixins import OrderableAggMixin
 from django.db.models import F, Func
 
 
@@ -10,10 +11,19 @@ class CollateAsChar(Func):
     template = '(%(expressions)s) COLLATE "%(function)s"'
 
 
-class EmptyGroupByJSONBAgg(JSONBAgg):
+class OrderableJSONBAgg(OrderableAggMixin, JSONBAgg):
+    """
+    TODO in Django 3.2 ordering is supported natively on JSONBAgg so this is no longer needed.
+    """
+    template = '%(function)s(%(distinct)s%(expressions)s %(ordering)s)'
+
+
+class EmptyGroupByJSONBAgg(OrderableJSONBAgg):
     """
     JSONBAgg is a builtin aggregation function which means it includes the use of a GROUP BY clause.
     When used as an annotation for collecting config context data objects, the GROUP BY is
     incorrect. This subclass overrides the Django ORM aggregation control to remove the GROUP BY.
+
+    TODO in Django 3.2 ordering is supported natively on JSONBAgg so we only need to inherit from JSONBAgg.
     """
     contains_aggregate = False