From bbea5d0fcab117aa70424848b0840d543dc6a595 Mon Sep 17 00:00:00 2001
From: Felix Peters <felix.peters@breuninger.de>
Date: Wed, 14 Sep 2022 19:58:54 +0200
Subject: [PATCH] feat: Implement custom fields Add custom fields to lables on
 IPs, VMs and devices

---
 netbox_prometheus_sd/api/serializers.py        |  8 +++-----
 netbox_prometheus_sd/api/utils.py              |  6 ++++++
 netbox_prometheus_sd/tests/test_serializers.py | 11 +++++++++--
 netbox_prometheus_sd/tests/utils.py            |  6 ++++++
 4 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/netbox_prometheus_sd/api/serializers.py b/netbox_prometheus_sd/api/serializers.py
index 9940086..281e7d7 100644
--- a/netbox_prometheus_sd/api/serializers.py
+++ b/netbox_prometheus_sd/api/serializers.py
@@ -34,6 +34,7 @@ def get_labels(self, obj):
         utils.extract_cluster(obj, labels)
         utils.extract_services(obj, labels)
         utils.extract_contacts(obj, labels)
+        utils.extract_custom_fields(obj, labels)
 
         if hasattr(obj, "device_role") and obj.device_role is not None:
             labels["role"] = obj.device_role.name
@@ -47,11 +48,6 @@ def get_labels(self, obj):
             labels["site"] = obj.site.name
             labels["site_slug"] = obj.site.slug
 
-        if hasattr(obj, "custom_field_data") and obj.custom_field_data is not None:
-            for key, value in obj.custom_field_data.items():
-                if value is not None:
-                    labels["custom_field_" + key.lower()] = value
-
         return labels.get_labels()
 
 
@@ -80,6 +76,7 @@ def get_labels(self, obj):
         utils.extract_cluster(obj, labels)
         utils.extract_services(obj, labels)
         utils.extract_contacts(obj, labels)
+        utils.extract_custom_fields(obj, labels)
 
         if hasattr(obj, "role") and obj.role is not None:
             labels["role"] = obj.role.name
@@ -121,5 +118,6 @@ def get_labels(self, obj):
 
         utils.extract_tags(obj, labels)
         utils.extract_tenant(obj, labels)
+        utils.extract_custom_fields(obj, labels)
 
         return labels.get_labels()
diff --git a/netbox_prometheus_sd/api/utils.py b/netbox_prometheus_sd/api/utils.py
index 029b545..4882547 100644
--- a/netbox_prometheus_sd/api/utils.py
+++ b/netbox_prometheus_sd/api/utils.py
@@ -88,3 +88,9 @@ def extract_contacts(obj, labels: LabelDict):
                 labels[f"contact_{contact.priority}_comments"] = contact.contact.comments
             if hasattr(contact, "role") and contact.role is not None:
                 labels[f"contact_{contact.priority}_role"] = contact.role.name
+
+def extract_custom_fields(obj, labels: LabelDict):
+    if hasattr(obj, "custom_field_data") and obj.custom_field_data is not None:
+        for key, value in obj.custom_field_data.items():
+            if value is not None:
+                labels["custom_field_" + key.lower()] = value
diff --git a/netbox_prometheus_sd/tests/test_serializers.py b/netbox_prometheus_sd/tests/test_serializers.py
index a9b8006..07f0088 100644
--- a/netbox_prometheus_sd/tests/test_serializers.py
+++ b/netbox_prometheus_sd/tests/test_serializers.py
@@ -71,8 +71,9 @@ def test_vm_full_to_target(self):
         self.assertDictContainsSubset(
             {"__meta_netbox_primary_ip6": "2001:db8:1701::2"}, data["labels"]
         )
-
-
+        self.assertDictContainsSubset(
+            {"__meta_netbox_custom_field_customer_id": "foobar-123"}, data["labels"]
+        )
 
 
 class PrometheusDeviceSerializerTests(TestCase):
@@ -129,6 +130,9 @@ def test_device_full_to_target(self):
         self.assertDictContainsSubset(
             {"__meta_netbox_tenant_slug": "acme"}, data["labels"]
         )
+        self.assertDictContainsSubset(
+            {"__meta_netbox_custom_field_customer_id": "foobar-123"}, data["labels"]
+        )
 
 
 class PrometheusIPAddressSerializerTests(TestCase):
@@ -180,3 +184,6 @@ def test_ip_full_to_target(self):
         self.assertDictContainsSubset(
             {"__meta_netbox_tenant_group_slug": "federation"}, data["labels"]
         )
+        self.assertDictContainsSubset(
+            {"__meta_netbox_custom_field_customer_id": "foobar-123"}, data["labels"]
+        )
diff --git a/netbox_prometheus_sd/tests/utils.py b/netbox_prometheus_sd/tests/utils.py
index 1b7c5b6..943c217 100644
--- a/netbox_prometheus_sd/tests/utils.py
+++ b/netbox_prometheus_sd/tests/utils.py
@@ -26,6 +26,9 @@ def build_tenant():
     return Tenant.objects.get_or_create(name="Acme Corp.", slug="acme")[0]
 
 
+def build_custom_fields():
+    return { "customer_id": "foobar-123"}
+
 def build_minimal_vm(name):
     return VirtualMachine.objects.get_or_create(name=name, cluster=build_cluster())[0]
 
@@ -33,6 +36,7 @@ def build_minimal_vm(name):
 def build_vm_full(name):
     vm = build_minimal_vm(name=name)
     vm.tenant = build_tenant()
+    vm.custom_field_data = build_custom_fields()
     vm.role = DeviceRole.objects.get_or_create(name="VM", slug="vm", vm_role=True)[0]
     vm.platform = Platform.objects.get_or_create(
         name="Ubuntu 20.04", slug="ubuntu-20.04"
@@ -65,6 +69,7 @@ def build_minimal_device(name):
 def build_device_full(name):
     device = build_minimal_device(name)
     device.tenant = build_tenant()
+    device.custom_field_data = build_custom_fields()
     device.platform = Platform.objects.get_or_create(name="Junos", slug="junos")[0]
     device.primary_ip4 = IPAddress.objects.get_or_create(address="192.168.0.1/24")[0]
     device.primary_ip6 = IPAddress.objects.get_or_create(address="2001:db8:1701::2/64")[
@@ -81,6 +86,7 @@ def build_minimal_ip(address):
 
 def build_full_ip(address, dns_name=""):
     ip = build_minimal_ip(address=address)
+    ip.custom_field_data = build_custom_fields()
     ip.tenant = Tenant.objects.get_or_create(
         name="Starfleet",
         slug="starfleet",