From 31c9f865c6a67efb5a05f08c94103ccec2ec6bd3 Mon Sep 17 00:00:00 2001
From: Brian Candler <b.candler@pobox.com>
Date: Thu, 31 Aug 2023 11:43:38 +0100
Subject: [PATCH] Use Device.role instead of Device.device_role where available
 (Netbox >=3.6.0)

Fixes #132

Signed-off-by: Brian Candler <b.candler@pobox.com>
---
 .github/workflows/ci.yml                | 1 +
 netbox_prometheus_sd/api/serializers.py | 5 ++++-
 netbox_prometheus_sd/api/views.py       | 2 +-
 netbox_prometheus_sd/tests/utils.py     | 7 ++++---
 4 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 1db23ba..b4a955d 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -19,6 +19,7 @@ jobs:
           - v3.3.2
           - v3.4.4
           - v3.5.3
+          - v3.6.0
           - master
 
     steps:
diff --git a/netbox_prometheus_sd/api/serializers.py b/netbox_prometheus_sd/api/serializers.py
index 0634d43..422ee49 100644
--- a/netbox_prometheus_sd/api/serializers.py
+++ b/netbox_prometheus_sd/api/serializers.py
@@ -82,7 +82,10 @@ def get_labels(self, obj):
         utils.extract_rack(obj, labels)
         utils.extract_custom_fields(obj, labels)
 
-        if hasattr(obj, "device_role") and obj.device_role is not None:
+        if hasattr(obj, "role") and obj.role is not None:
+            labels["role"] = obj.role.name
+            labels["role_slug"] = obj.role.slug
+        elif hasattr(obj, "device_role") and obj.device_role is not None:  # netbox <3.6.0
             labels["role"] = obj.device_role.name
             labels["role_slug"] = obj.device_role.slug
 
diff --git a/netbox_prometheus_sd/api/views.py b/netbox_prometheus_sd/api/views.py
index fc67c03..5903644 100644
--- a/netbox_prometheus_sd/api/views.py
+++ b/netbox_prometheus_sd/api/views.py
@@ -75,7 +75,7 @@ class VirtualMachineViewSet(
 class DeviceViewSet(NetboxPrometheusSDModelViewSet):  # pylint: disable=too-many-ancestors
     queryset = Device.objects.prefetch_related(
         "device_type__manufacturer",
-        "device_role",
+        "role" if hasattr(Device, "role") else "device_role",
         "tenant",
         "platform",
         "site",
diff --git a/netbox_prometheus_sd/tests/utils.py b/netbox_prometheus_sd/tests/utils.py
index 9d5d3e5..a33f246 100644
--- a/netbox_prometheus_sd/tests/utils.py
+++ b/netbox_prometheus_sd/tests/utils.py
@@ -92,17 +92,18 @@ def build_vm_full(name):
 
 
 def build_minimal_device(name):
+    role_attr = "role" if hasattr(Device, "role") else "device_role"
     return Device.objects.get_or_create(
         name=name,
-        device_role=DeviceRole.objects.get_or_create(name="Firewall", slug="firewall")[
-            0
-        ],
         device_type=DeviceType.objects.get_or_create(
             model="SRX",
             slug="srx",
             manufacturer=Manufacturer.objects.get_or_create(
                 name="Juniper", slug="juniper"
             )[0],
+        **{
+            role_attr: DeviceRole.objects.get_or_create(name="Firewall", slug="firewall")[0],
+        }
         )[0],
         site=Site.objects.get_or_create(name="Site", slug="site")[0],
     )[0]