Closes #38: Add role field to ConfigDiffScript

docs/colliecting-diffs.md

@@ -28,8 +28,8 @@ You can find scripts list in navbar `Customization -> Scripts`.
 
 ![Screenshot of the scripts list](media/screenshots/script-list.png)
 
-In the script, you can define a site, on which devices run compliance, or devices.
- If you define both fields, script will run only on devices from `Devices` field
+In the script, you can define a site or role, on which devices run compliance, or devices.
+ If you define all fields, script will run only on devices from `Devices` field
 
 !!! warning
     Script runs only on devices with assigned Primary IP, Platform and PlatformSetting

netbox_config_diff/compliance/base.py

@@ -6,7 +6,8 @@
 from core.choices import DataSourceStatusChoices
 from core.models import DataFile, DataSource
 from dcim.choices import DeviceStatusChoices
-from dcim.models import Device, Site
+from dcim.models import Device, DeviceRole, Site
+from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
 from django.db.models import Q
 from extras.scripts import MultiObjectVar, ObjectVar
@@ -27,14 +28,19 @@ class ConfigDiffBase(SecretsMixin):
         required=False,
         description="Run compliance for devices (with primary IP, platform) in this site",
     )
+    role = ObjectVar(
+        model=DeviceRole,
+        required=False,
+        description="Run compliance for devices with this role",
+    )
     devices = MultiObjectVar(
         model=Device,
         required=False,
         query_params={
             "has_primary_ip": True,
             "platform_id__n": "null",
         },
-        description="If you define devices in this field, the Site field will be ignored",
+        description="If you define devices in this field, Site, Role fields will be ignored",
     )
     status = CustomChoiceVar(
         choices=DeviceStatusChoices,
@@ -57,8 +63,8 @@ def run_script(self, data: dict) -> None:
         self.update_in_db(devices)
 
     def validate_data(self, data: dict) -> Iterable[Device]:
-        if not data["site"] and not data["devices"]:
-            raise AbortScript("Define site or devices")
+        if not data["site"] and not data["role"] and not data["devices"]:
+            raise AbortScript("Define site, role or devices")
         if data.get("data_source") and data["data_source"].status != DataSourceStatusChoices.COMPLETED:
             raise AbortScript("Define synced DataSource")
 
@@ -75,13 +81,21 @@ def validate_data(self, data: dict) -> Iterable[Device]:
                 )
             )
         else:
-            devices = Device.objects.filter(
-                site=data["site"],
-                status=data["status"],
-                platform__platform_setting__isnull=False,
-            ).exclude(
+            filters = {
+                "status": data["status"],
+                "platform__platform_setting__isnull": False,
+            }
+            if data["site"]:
+                filters["site"] = data["site"]
+            elif data["role"]:
+                if settings.VERSION.split(".", 1)[1].startswith("5"):
+                    filters["device_role"] = data["role"]
+                else:
+                    filters["role"] = data["role"]
+            devices = Device.objects.filter(**filters).exclude(
                 Q(primary_ip4__isnull=True) & Q(primary_ip6__isnull=True),
             )
+
         if data.get("devices"):
             if qs_diff := data["devices"].difference(devices):
                 platforms = {d.platform.name for d in qs_diff}

tests/conftest.py

@@ -34,6 +34,7 @@ def script_data_factory() -> "ScriptDataFactory":
     def factory(**fields: Unpack["ScriptData"]) -> "ScriptData":
         data = {
             "site": None,
+            "role": None,
             "devices": None,
             "data_source": None,
             "status": "active",

tests/test_compliance.py

@@ -16,7 +16,7 @@
 def test_validate_data_no_data(mock_config_diff: "ConfigDiffBase", script_data: "ScriptData") -> None:
     with pytest.raises(AbortScript) as e:
         mock_config_diff.validate_data(data=script_data)
-    assert str(e.value) == "Define site or devices"
+    assert str(e.value) == "Define site, role or devices"
 
 
 @pytest.mark.django_db()