Add site and site_slug labels to device services

FlxPeters · Apr 4, 2024 · 0a2f612 · 0a2f612
1 parent b40f30e
commit 0a2f612
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 5 deletions.
diff --git a/netbox_prometheus_sd/api/serializers.py b/netbox_prometheus_sd/api/serializers.py
@@ -96,10 +96,6 @@ def get_labels(self, obj):
             labels["device_type"] = obj.device_type.model
             labels["device_type_slug"] = obj.device_type.slug
 
-        if hasattr(obj, "site") and obj.site is not None:
-            labels["site"] = obj.site.name
-            labels["site_slug"] = obj.site.slug
-
         labels = labels.get_labels()
 
         # Those shouldn't have the netbox prefix

diff --git a/netbox_prometheus_sd/api/utils.py b/netbox_prometheus_sd/api/utils.py
@@ -62,6 +62,11 @@ def extract_cluster(obj, labels: LabelDict):
             labels["site"] = obj.cluster.site.name
             labels["site_slug"] = obj.cluster.site.slug
 
+    # Has precedence over cluster site
+    if hasattr(obj, "site") and obj.site is not None:
+        labels["site"] = obj.site.name
+        labels["site_slug"] = obj.site.slug
+
 
 def extract_primary_ip(obj, labels: LabelDict):
     if getattr(obj, "primary_ip", None) is not None:

diff --git a/netbox_prometheus_sd/tests/test_serializers.py b/netbox_prometheus_sd/tests/test_serializers.py
@@ -191,6 +191,12 @@ def test_device_full_to_target(self):
             {"__meta_netbox_oob_ip": "10.0.0.1"}, data["labels"]
         )
         self.assertDictContainsSubset({"__meta_netbox_rack": "R01B01"}, data["labels"])
+        self.assertDictContainsSubset(
+            {"__meta_netbox_site": "Site"}, data["labels"]
+        )
+        self.assertDictContainsSubset(
+            {"__meta_netbox_site_slug": "site"}, data["labels"]
+        )
         self.assertDictContainsSubset(
             {"__meta_netbox_description": "Device Description"}, data["labels"]
         )
@@ -266,7 +272,46 @@ def test_ip_full_to_target(self):
 
 
 class PrometheusServiceSerializerTests(TestCase):
-    def test_service_full_to_target(self):
+    def test_device_service_full_to_target(self):
+        device = utils.build_device_full("firewall-full-01")
+        instance = device.services.first()
+        data_list = PrometheusServiceSerializer(many=True, instance=[instance]).data
+
+        self.assertEqual(data_list[0]["targets"], ["ssh"])
+        for data in data_list:
+            self.assertDictContainsSubset({"__meta_netbox_id": str(instance.id)}, data["labels"])
+            self.assertDictContainsSubset(
+                {"__meta_netbox_display": "ssh (TCP/22)"}, data["labels"]
+            )
+            self.assertDictContainsSubset(
+                {"__meta_netbox_ports": "22"}, data["labels"]
+            )
+            self.assertDictContainsSubset(
+                {"__meta_netbox_parent": "firewall-full-01"}, data["labels"]
+            )
+            self.assertDictContainsSubset(
+                {"__meta_netbox_tenant": "Acme Corp."}, data["labels"]
+            )
+            self.assertDictContainsSubset(
+                {"__meta_netbox_tenant_slug": "acme"}, data["labels"]
+            )
+            self.assertDictContainsSubset(
+                {"__meta_netbox_site": "Site"}, data["labels"]
+            )
+            self.assertDictContainsSubset(
+                {"__meta_netbox_site_slug": "site"}, data["labels"]
+            )
+            self.assertDictContainsSubset(
+                {"__meta_netbox_primary_ip": "2001:db8:1701::2"}, data["labels"]
+            )
+            self.assertDictContainsSubset(
+                {"__meta_netbox_primary_ip4": "192.168.0.1"}, data["labels"]
+            )
+            self.assertDictContainsSubset(
+                {"__meta_netbox_primary_ip6": "2001:db8:1701::2"}, data["labels"]
+            )
+
+    def test_vm_service_full_to_target(self):
         vm = utils.build_vm_full("vm-full-01.example.com")
         instance = vm.services.first()
         data_list = PrometheusServiceSerializer(many=True, instance=[instance]).data

diff --git a/netbox_prometheus_sd/tests/utils.py b/netbox_prometheus_sd/tests/utils.py
@@ -189,8 +189,11 @@ def build_device_full(name):
     device.rack = Rack.objects.get_or_create(
         name="R01B01", site=Site.objects.get_or_create(name="Site", slug="site")[0]
     )[0]
+    device.site = Site.objects.get_or_create(name="Site", slug="site")[0]
     device.tags.add("Tag1")
     device.tags.add("Tag 2")
+
+    Service.objects.create(device=device, name="ssh", protocol='tcp', ports=[22])
     return device