Issue #14962 VM to merge directly related site context (#14992)

* Issue #14962 VM to merge directly related site context

* Cleanup & rewrite test

---------

Co-authored-by: Chris Carter <chris.carter@spinlocksecurity.com>
Co-authored-by: Jeremy Stretch <jstretch@netboxlabs.com>

netbox/extras/querysets.py

@@ -120,34 +120,29 @@ class ConfigContextModelQuerySet(RestrictedQuerySet):
         if self.model._meta.model_name == 'device':
             base_query.add((Q(locations=OuterRef('location')) | Q(locations=None)), Q.AND)
             base_query.add((Q(device_types=OuterRef('device_type')) | Q(device_types=None)), Q.AND)
-            base_query.add((Q(roles=OuterRef('role')) | Q(roles=None)), Q.AND)
-            base_query.add((Q(sites=OuterRef('site')) | Q(sites=None)), Q.AND)
-            region_field = 'site__region'
-            sitegroup_field = 'site__group'
 
         elif self.model._meta.model_name == 'virtualmachine':
-            base_query.add((Q(roles=OuterRef('role')) | Q(roles=None)), Q.AND)
-            base_query.add((Q(sites=OuterRef('cluster__site')) | Q(sites=None)), Q.AND)
             base_query.add(Q(device_types=None), Q.AND)
-            region_field = 'cluster__site__region'
-            sitegroup_field = 'cluster__site__group'
+
+        base_query.add((Q(roles=OuterRef('role')) | Q(roles=None)), Q.AND)
+        base_query.add((Q(sites=OuterRef('site')) | Q(sites=None)), Q.AND)
 
         base_query.add(
             (Q(
-                regions__tree_id=OuterRef(f'{region_field}__tree_id'),
-                regions__level__lte=OuterRef(f'{region_field}__level'),
-                regions__lft__lte=OuterRef(f'{region_field}__lft'),
-                regions__rght__gte=OuterRef(f'{region_field}__rght'),
+                regions__tree_id=OuterRef('site__region__tree_id'),
+                regions__level__lte=OuterRef('site__region__level'),
+                regions__lft__lte=OuterRef('site__region__lft'),
+                regions__rght__gte=OuterRef('site__region__rght'),
             ) | Q(regions=None)),
             Q.AND
         )
 
         base_query.add(
             (Q(
-                site_groups__tree_id=OuterRef(f'{sitegroup_field}__tree_id'),
-                site_groups__level__lte=OuterRef(f'{sitegroup_field}__level'),
-                site_groups__lft__lte=OuterRef(f'{sitegroup_field}__lft'),
-                site_groups__rght__gte=OuterRef(f'{sitegroup_field}__rght'),
+                site_groups__tree_id=OuterRef('site__group__tree_id'),
+                site_groups__level__lte=OuterRef('site__group__level'),
+                site_groups__lft__lte=OuterRef('site__group__lft'),
+                site_groups__rght__gte=OuterRef('site__group__rght'),
             ) | Q(site_groups=None)),
             Q.AND
         )

netbox/extras/tests/test_models.py

@@ -270,7 +270,12 @@ class ConfigContextTest(TestCase):
         tag = Tag.objects.first()
         cluster_type = ClusterType.objects.create(name="Cluster Type")
         cluster_group = ClusterGroup.objects.create(name="Cluster Group")
-        cluster = Cluster.objects.create(name="Cluster", group=cluster_group, type=cluster_type)
+        cluster = Cluster.objects.create(
+            name="Cluster",
+            group=cluster_group,
+            type=cluster_type,
+            site=site,
+        )
 
         region_context = ConfigContext.objects.create(
             name="region",
@@ -354,6 +359,41 @@ class ConfigContextTest(TestCase):
         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())
 
+    def test_virtualmachine_site_context(self):
+        """
+        Check that config context associated with a site applies to a VM whether the VM is assigned
+        directly to that site or via its cluster.
+        """
+        site = Site.objects.first()
+        cluster_type = ClusterType.objects.create(name="Cluster Type")
+        cluster = Cluster.objects.create(name="Cluster", type=cluster_type, site=site)
+        vm_role = DeviceRole.objects.first()
+
+        # Create a ConfigContext associated with the site
+        context = ConfigContext.objects.create(
+            name="context1",
+            weight=100,
+            data={"foo": True}
+        )
+        context.sites.add(site)
+
+        # Create one VM assigned directly to the site, and one assigned via the cluster
+        vm1 = VirtualMachine.objects.create(name="VM 1", site=site, role=vm_role)
+        vm2 = VirtualMachine.objects.create(name="VM 2", cluster=cluster, role=vm_role)
+
+        # Check that their individually-rendered config contexts are identical
+        self.assertEqual(
+            vm1.get_config_context(),
+            vm2.get_config_context()
+        )
+
+        # Check that their annotated config contexts are identical
+        vms = VirtualMachine.objects.filter(pk__in=(vm1.pk, vm2.pk)).annotate_config_context_data()
+        self.assertEqual(
+            vms[0].get_config_context(),
+            vms[1].get_config_context()
+        )
+
     def test_multiple_tags_return_distinct_objects(self):
         """
         Tagged items use a generic relationship, which results in duplicate rows being returned when queried.