Rename CustomFieldModelViewSet name for NetBox 3.2+ (#50)

* Fix ModelViewSet name for NetBox 3.2+ * Add NetBox v3.2.2 as new test target. * Ensure imports work for Netbox <3.2 and >=3.2.
FlxPeters · May 9, 2022 · 1545b11 · 1545b11
1 parent c64408c
commit 1545b11
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -18,8 +18,8 @@ jobs:
           - v2.11.11
           - v3.0.12
           - v3.1.8
-          # Remove master for now (see https://github.com/FlxPeters/netbox-plugin-prometheus-sd/issues/46)
-          # - master
+          - v3.2.2
+          - master
 
     steps:
       - uses: actions/checkout@v2

diff --git a/netbox_prometheus_sd/api/views.py b/netbox_prometheus_sd/api/views.py
@@ -1,7 +1,14 @@
 from ipam.models import IPAddress
 from virtualization.models import VirtualMachine
 from dcim.models.devices import Device
-from extras.api.views import CustomFieldModelViewSet
+
+# The base ViewSet has been renamed, this try-except helps to support
+# Both < 3.2 and the newer 3.2+ Versions:
+# https://github.com/netbox-community/netbox/commit/bbdeae0ed9bcc06fb96ffa2970272e1a3447448c
+try:
+    from netbox.api.viewsets import NetBoxModelViewSet
+except ImportError:
+    from extras.api.views import CustomFieldModelViewSet as NetBoxModelViewSet
 
 # Filtersets have been renamed, we support both
 # https://github.com/netbox-community/netbox/commit/1024782b9e0abb48f6da65f8248741227d53dbed#diff-d9224204dab475bbe888868c02235b8ef10f07c9201c45c90804d395dc161c40
@@ -25,7 +32,7 @@
 
 
 class VirtualMachineViewSet(
-    CustomFieldModelViewSet
+    NetBoxModelViewSet
 ):  # pylint: disable=too-many-ancestors
     queryset = VirtualMachine.objects.prefetch_related(
         "cluster__site",
@@ -44,7 +51,7 @@ class VirtualMachineViewSet(
     pagination_class = None
 
 
-class DeviceViewSet(CustomFieldModelViewSet):  # pylint: disable=too-many-ancestors
+class DeviceViewSet(NetBoxModelViewSet):  # pylint: disable=too-many-ancestors
     queryset = Device.objects.prefetch_related(
         "device_type__manufacturer",
         "device_role",
@@ -64,7 +71,7 @@ class DeviceViewSet(CustomFieldModelViewSet):  # pylint: disable=too-many-ancest
     pagination_class = None
 
 
-class IPAddressViewSet(CustomFieldModelViewSet):  # pylint: disable=too-many-ancestors
+class IPAddressViewSet(NetBoxModelViewSet):  # pylint: disable=too-many-ancestors
     queryset = IPAddress.objects.prefetch_related("tenant", "tags")
     serializer_class = PrometheusIPAddressSerializer
     filterset_class = IPAddressFilterSet