Merge pull request #10 from nscaledev/sr/add_cluster_lookup_devices

steve-nscale · web-flow · commit ba63a0070880 · 2026-01-09T15:29:54.000Z
feat: add cluster name filter for devices and fix cluster scope labels
diff --git a/netbox_prometheus_sd/api/utils.py b/netbox_prometheus_sd/api/utils.py
@@ -66,21 +66,23 @@ def extract_cluster(obj, labels: LabelDict):
             labels["cluster_group"] = obj.cluster.group.name
         if obj.cluster.type:
             labels["cluster_type"] = obj.cluster.type.name
-        try: # Netbox >4.2
-            if obj.cluster.scope:
-                labels["scope"] = obj.cluster.scope.name
-                labels["scope_slug"] = obj.cluster.scope.slug
-        except AttributeError: # Netbox <4.2
-            if obj.cluster.site:
-                labels["site"] = obj.cluster.site.name
-                labels["site_slug"] = obj.cluster.site.slug
-
-    # Has precedence over cluster scope
-    if hasattr(obj, "scope") and obj.scope is not None:
-        labels["scope"] = obj.scope.name
-        labels["scope_slug"] = obj.scope.slug
-
-    # Still Return site labels for Devices
+        # NetBox 4.2+ uses scope (generic FK) instead of site
+        if hasattr(obj.cluster, "scope") and obj.cluster.scope is not None:
+            scope = obj.cluster.scope
+            # scope can be region, site_group, site, or location
+            if hasattr(scope, "slug"):
+                labels["cluster_scope"] = scope.name
+                labels["cluster_scope_slug"] = scope.slug
+            # If scope is a site, also set site labels
+            if scope.__class__.__name__ == "Site":
+                labels["site"] = scope.name
+                labels["site_slug"] = scope.slug
+        # NetBox < 4.2 uses site directly
+        elif hasattr(obj.cluster, "site") and obj.cluster.site:
+            labels["site"] = obj.cluster.site.name
+            labels["site_slug"] = obj.cluster.site.slug
+
+    # Has precedence over cluster site/scope
     if hasattr(obj, "site") and obj.site is not None:
         labels["site"] = obj.site.name
         labels["site_slug"] = obj.site.slug
diff --git a/netbox_prometheus_sd/api/views.py b/netbox_prometheus_sd/api/views.py
@@ -30,15 +30,13 @@ class NetboxPrometheusSDModelViewSet(
 # https://github.com/netbox-community/netbox/commit/1024782b9e0abb48f6da65f8248741227d53dbed#diff-d9224204dab475bbe888868c02235b8ef10f07c9201c45c90804d395dc161c40
 try:
     from ipam.filtersets import IPAddressFilterSet
-    from dcim.filtersets import DeviceFilterSet
     from virtualization.filtersets import VirtualMachineFilterSet
 except ImportError:
     from ipam.filters import IPAddressFilterSet
-    from dcim.filters import DeviceFilterSet
     from virtualization.filters import VirtualMachineFilterSet
 
 
-from ..filtersets import ServiceFilterSet
+from ..filtersets import ServiceFilterSet, DeviceFilterSet
 from .serializers import (
     PrometheusIPAddressSerializer,
     PrometheusDeviceSerializer,
@@ -94,6 +92,8 @@ class DeviceViewSet(NetboxPrometheusSDModelViewSet):
         "primary_ip4__nat_outside",
         "primary_ip6__nat_outside",
         "tags",
+        "cluster__group",
+        "cluster__type",
     )
     filterset_class = DeviceFilterSet
     serializer_class = PrometheusDeviceSerializer
diff --git a/netbox_prometheus_sd/filtersets.py b/netbox_prometheus_sd/filtersets.py
@@ -14,6 +14,11 @@
 except ImportError:
     from ipam.filters import ServiceFilterSet as NetboxServiceFilterSet
 
+try:
+    from dcim.filtersets import DeviceFilterSet as NetboxDeviceFilterSet
+except ImportError:
+    from dcim.filters import DeviceFilterSet as NetboxDeviceFilterSet
+
 
 class ServiceFilterSet(NetboxServiceFilterSet):
     """Filter set to support tenancy over the device/VM foreign key.
@@ -59,3 +64,23 @@ def filter_by_tenant_slug(self, queryset, name, value):
             Q(device__tenant__slug__in=value)
             | Q(virtual_machine__tenant__slug__in=value)
         )
+
+
+class DeviceFilterSet(NetboxDeviceFilterSet):
+    """Extends NetBox's DeviceFilterSet to add cluster name filter.
+
+    NetBox's DeviceFilterSet only has cluster_id, not cluster (name).
+    This adds the missing cluster name filter for consistency with other filters.
+    Cluster model doesn't have a slug field, so we filter by name (case-insensitive).
+    """
+
+    cluster = MultiValueCharFilter(
+        method='filter_by_cluster_name',
+        label=_('Cluster (name)'),
+    )
+
+    def filter_by_cluster_name(self, queryset, name, value):
+        q = Q()
+        for v in value:
+            q |= Q(cluster__name__iexact=v)
+        return queryset.filter(q)
diff --git a/netbox_prometheus_sd/tests/test_filtersets.py b/netbox_prometheus_sd/tests/test_filtersets.py
@@ -1,11 +1,12 @@
 from django.test import TestCase
 
+from dcim.models import Device
 from ipam.models import Service
 from tenancy.models import Tenant
 from utilities.testing import ChangeLoggedFilterSetTests
 
 from . import utils
-from ..filtersets import ServiceFilterSet
+from ..filtersets import ServiceFilterSet, DeviceFilterSet
 
 
 class ServiceTestCase(TestCase, ChangeLoggedFilterSetTests):
@@ -34,3 +35,47 @@ def test_vm_tenant(self):
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 6)
         params = {"tenant": [tenant.slug]}
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 6)
+
+
+class DeviceFilterSetTestCase(TestCase):
+    """Test cases for DeviceFilterSet cluster name filter.
+
+    Note: We don't inherit ChangeLoggedFilterSetTests because NetBox's own
+    DeviceFilterSet is missing some filters (vc_master_for_id) which would
+    cause that test to fail.
+    """
+    queryset = Device.objects.all()
+    filterset = DeviceFilterSet
+
+    @classmethod
+    def setUpTestData(cls):
+        """Create test devices with and without clusters."""
+        # Devices with clusters
+        utils.build_device_with_cluster("device-cluster-01", cluster_name="SYS2-STA1")
+        utils.build_device_with_cluster("device-cluster-02", cluster_name="SYS2-STA1")
+        utils.build_device_with_cluster("device-cluster-03", cluster_name="SYS3-STA2")
+        # Device without cluster
+        utils.build_minimal_device("device-no-cluster-01")
+
+    def test_filter_by_cluster_name_exact(self):
+        """Test filtering devices by exact cluster name."""
+        params = {'cluster': ['SYS2-STA1']}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+
+    def test_filter_by_cluster_name_case_insensitive(self):
+        """Test filtering devices by cluster name is case-insensitive."""
+        params = {'cluster': ['sys2-sta1']}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+
+        params = {'cluster': ['Sys2-Sta1']}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+
+    def test_filter_by_cluster_name_multiple(self):
+        """Test filtering devices by multiple cluster names."""
+        params = {'cluster': ['SYS2-STA1', 'SYS3-STA2']}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 3)
+
+    def test_filter_by_cluster_name_no_match(self):
+        """Test filtering devices by non-existent cluster name returns empty."""
+        params = {'cluster': ['NONEXISTENT']}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 0)
diff --git a/netbox_prometheus_sd/tests/test_serializers.py b/netbox_prometheus_sd/tests/test_serializers.py
@@ -1,15 +1,28 @@
 from django.test import TestCase
+from unittest.mock import MagicMock
+
 from ..api.serializers import (
     PrometheusDeviceSerializer,
     PrometheusIPAddressSerializer,
     PrometheusServiceSerializer,
     PrometheusVirtualMachineSerializer,
 )
+from ..api.utils import LabelDict, extract_cluster
 from . import utils
 
 from ..api.utils import NETBOX_RELEASE_CURRENT, NETBOX_RELEASE_41
 
-class PrometheusVirtualMachineSerializerTests(TestCase):
+class DictSubsetMixin:
+    """Mixin to provide assertDictContainsSubset which was removed in Python 3.12."""
+
+    def assertDictContainsSubset(self, subset, dictionary, msg=None):
+        """Check that all key/value pairs in subset are in dictionary."""
+        for key, value in subset.items():
+            self.assertIn(key, dictionary, msg=msg)
+            self.assertEqual(dictionary[key], value, msg=msg)
+
+
+class PrometheusVirtualMachineSerializerTests(DictSubsetMixin, TestCase):
     def test_vm_minimal_to_target(self):
 
         instance = utils.build_minimal_vm("vm-01.example.com")
@@ -102,23 +115,12 @@ def test_vm_full_to_target(self):
             if NETBOX_RELEASE_CURRENT > NETBOX_RELEASE_41:
                 self.assertTrue(
                     utils.dictContainsSubset(
-                        {"__meta_netbox_scope": "Campus A"}, data["labels"]
-                    )
-                )
-                self.assertTrue(
-                    utils.dictContainsSubset(
-                        {"__meta_netbox_scope_slug": "campus-a"}, data["labels"]
-                    )
-                )
-            else:
-                self.assertTrue(
-                    utils.dictContainsSubset(
-                        {"__meta_netbox_site": "Campus A"}, data["labels"]
+                        {"__meta_netbox_cluster_scope": "Campus A"}, data["labels"]
                     )
                 )
                 self.assertTrue(
                     utils.dictContainsSubset(
-                        {"__meta_netbox_site_slug": "campus-a"}, data["labels"]
+                        {"__meta_netbox_cluster_scope_slug": "campus-a"}, data["labels"]
                     )
                 )
             self.assertTrue(
@@ -200,7 +202,7 @@ def test_vm_full_to_target(self):
             )
 
 
-class PrometheusDeviceSerializerTests(TestCase):
+class PrometheusDeviceSerializerTests(DictSubsetMixin, TestCase):
     def test_device_minimal_to_target(self):
         instance = utils.build_minimal_device("firewall-01")
         data = PrometheusDeviceSerializer(many=True, instance=[instance]).data[0]
@@ -364,7 +366,7 @@ def test_device_full_to_target(self):
         )
 
 
-class PrometheusIPAddressSerializerTests(TestCase):
+class PrometheusIPAddressSerializerTests(DictSubsetMixin, TestCase):
     def test_ip_minimal_to_target(self):
         instance = utils.build_minimal_ip("10.10.10.10/24")
         data = PrometheusIPAddressSerializer(many=True, instance=[instance]).data[0]
@@ -440,7 +442,7 @@ def test_ip_full_to_target(self):
         )
 
 
-class PrometheusServiceSerializerTests(TestCase):
+class PrometheusServiceSerializerTests(DictSubsetMixin, TestCase):
     def test_device_service_full_to_target(self):
         device = utils.build_device_full("firewall-full-01")
         instance = device.services.first()
@@ -548,23 +550,12 @@ def test_vm_service_full_to_target(self):
             if NETBOX_RELEASE_CURRENT > NETBOX_RELEASE_41:
                 self.assertTrue(
                     utils.dictContainsSubset(
-                        {"__meta_netbox_scope": "Campus A"}, data["labels"]
+                        {"__meta_netbox_cluster_scope": "Campus A"}, data["labels"]
                     )
                 )
                 self.assertTrue(
                     utils.dictContainsSubset(
-                        {"__meta_netbox_scope_slug": "campus-a"}, data["labels"]
-                    )
-                )
-            else:
-                self.assertTrue(
-                    utils.dictContainsSubset(
-                        {"__meta_netbox_site": "Campus A"}, data["labels"]
-                    )
-                )
-                self.assertTrue(
-                    utils.dictContainsSubset(
-                        {"__meta_netbox_site_slug": "campus-a"}, data["labels"]
+                        {"__meta_netbox_cluster_scope_slug": "campus-a"}, data["labels"]
                     )
                 )
             self.assertTrue(
@@ -582,3 +573,88 @@ def test_vm_service_full_to_target(self):
                     {"__meta_netbox_primary_ip6": "2001:db8:1701::2"}, data["labels"]
                 )
             )
+
+
+class ExtractClusterTests(TestCase):
+    """Tests for extract_cluster function handling both NetBox 4.2+ scope and older site."""
+
+    def test_extract_cluster_basic_info(self):
+        """Test that basic cluster info (name, group, type) is extracted."""
+        device = utils.build_device_with_cluster("device-cluster-test", cluster_name="TestCluster")
+        labels = LabelDict()
+        extract_cluster(device, labels)
+
+        self.assertEqual(labels.get("cluster"), "TestCluster")
+        self.assertEqual(labels.get("cluster_group"), "VMware")
+        self.assertEqual(labels.get("cluster_type"), "On Prem")
+
+    def test_extract_cluster_scope_from_scope(self):
+        """Test that scope labels are extracted from cluster scope (NetBox 4.2+)."""
+        device = utils.build_device_with_cluster("device-scope-test", cluster_name="ScopeCluster")
+        labels = LabelDict()
+        extract_cluster(device, labels)
+
+        # In NetBox 4.2+, cluster scope is labeled as 'cluster_scope'
+        self.assertEqual(labels.get("cluster_scope"), "Cluster Site")
+        self.assertEqual(labels.get("cluster_scope_slug"), "cluster-site")
+        # Device's own site takes precedence for 'site' label
+        self.assertEqual(labels.get("site"), "Site")
+        self.assertEqual(labels.get("site_slug"), "site")
+
+    def test_extract_cluster_no_cluster(self):
+        """Test that no cluster labels are added when device has no cluster."""
+        device = utils.build_minimal_device("device-no-cluster-test")
+        labels = LabelDict()
+        extract_cluster(device, labels)
+
+        self.assertIsNone(labels.get("cluster"))
+        self.assertIsNone(labels.get("cluster_group"))
+        self.assertIsNone(labels.get("cluster_type"))
+
+    def test_extract_cluster_device_site_overrides_cluster_site(self):
+        """Test that device's own site takes precedence over cluster site."""
+        device = utils.build_device_with_cluster("device-site-override", cluster_name="OverrideCluster")
+        # Device's own site should override cluster site
+        labels = LabelDict()
+        extract_cluster(device, labels)
+
+        # Device has its own site set in build_minimal_device ("Site", "site")
+        # which should take precedence over the cluster's site
+        self.assertEqual(labels.get("site"), "Site")
+        self.assertEqual(labels.get("site_slug"), "site")
+
+    def test_extract_cluster_with_scope_mock(self):
+        """Test extract_cluster with mocked scope object (NetBox 4.2+ style)."""
+        # Create a mock object to simulate NetBox 4.2+ cluster with scope
+        mock_site = MagicMock()
+        mock_site.name = "Scoped Site"
+        mock_site.slug = "scoped-site"
+        mock_site.__class__.__name__ = "Site"
+
+        mock_cluster = MagicMock()
+        mock_cluster.name = "MockCluster"
+        mock_cluster.group = MagicMock(name="MockGroup")
+        mock_cluster.group.name = "Mock Group"
+        mock_cluster.type = MagicMock(name="MockType")
+        mock_cluster.type.name = "Mock Type"
+        mock_cluster.scope = mock_site
+        # Ensure hasattr returns True for scope
+        del mock_cluster.site
+
+        mock_obj = MagicMock()
+        mock_obj.cluster = mock_cluster
+        # Remove site attribute from obj
+        del mock_obj.site
+
+        labels = LabelDict()
+        extract_cluster(mock_obj, labels)
+
+        self.assertEqual(labels.get("cluster"), "MockCluster")
+        self.assertEqual(labels.get("cluster_group"), "Mock Group")
+        self.assertEqual(labels.get("cluster_type"), "Mock Type")
+        # In NetBox 4.2+, cluster scope is labeled as 'cluster_scope'
+        self.assertEqual(labels.get("cluster_scope"), "Scoped Site")
+        self.assertEqual(labels.get("cluster_scope_slug"), "scoped-site")
+        # Since scope is a Site, site labels should also be set
+        self.assertEqual(labels.get("site"), "Scoped Site")
+        self.assertEqual(labels.get("site_slug"), "scoped-site")
diff --git a/netbox_prometheus_sd/tests/utils.py b/netbox_prometheus_sd/tests/utils.py
diff --git a/pyproject.toml b/pyproject.toml
