Add primary_ip4 and primary_ip6 filters for VirtualMachine and Virtua…

…lDeviceContext filtersets (netbox-community#13936)

netbox/dcim/filtersets.py

@@ -4,6 +4,7 @@
 
 from extras.filtersets import LocalConfigContextFilterSet
 from extras.models import ConfigTemplate
+from ipam.filtersets import PrimaryIPFilterSet
 from ipam.models import ASN, L2VPN, IPAddress, VRF
 from netbox.filtersets import (
     BaseFilterSet, ChangeLoggedModelFilterSet, OrganizationalModelFilterSet, NetBoxModelFilterSet,
@@ -817,7 +818,13 @@ class Meta:
         fields = ['id', 'name', 'slug', 'description']
 
 
-class DeviceFilterSet(NetBoxModelFilterSet, TenancyFilterSet, ContactModelFilterSet, LocalConfigContextFilterSet):
+class DeviceFilterSet(
+    NetBoxModelFilterSet,
+    TenancyFilterSet,
+    ContactModelFilterSet,
+    LocalConfigContextFilterSet,
+    PrimaryIPFilterSet,
+):
     manufacturer_id = django_filters.ModelMultipleChoiceFilter(
         field_name='device_type__manufacturer',
         queryset=Manufacturer.objects.all(),
@@ -993,16 +1000,6 @@ class DeviceFilterSet(NetBoxModelFilterSet, TenancyFilterSet, ContactModelFilter
         method='_device_bays',
         label=_('Has device bays'),
     )
-    primary_ip4_id = django_filters.ModelMultipleChoiceFilter(
-        field_name='primary_ip4',
-        queryset=IPAddress.objects.all(),
-        label=_('Primary IPv4 (ID)'),
-    )
-    primary_ip6_id = django_filters.ModelMultipleChoiceFilter(
-        field_name='primary_ip6',
-        queryset=IPAddress.objects.all(),
-        label=_('Primary IPv6 (ID)'),
-    )
     oob_ip_id = django_filters.ModelMultipleChoiceFilter(
         field_name='oob_ip',
         queryset=IPAddress.objects.all(),
@@ -1069,7 +1066,7 @@ def _device_bays(self, queryset, name, value):
         return queryset.exclude(devicebays__isnull=value)
 
 
-class VirtualDeviceContextFilterSet(NetBoxModelFilterSet, TenancyFilterSet):
+class VirtualDeviceContextFilterSet(NetBoxModelFilterSet, TenancyFilterSet, PrimaryIPFilterSet):
     device_id = django_filters.ModelMultipleChoiceFilter(
         field_name='device',
         queryset=Device.objects.all(),

netbox/dcim/tests/test_filtersets.py

@@ -4712,12 +4712,18 @@ def setUpTestData(cls):
         addresses = (
             IPAddress(assigned_object=interfaces[0], address='10.1.1.1/24'),
             IPAddress(assigned_object=interfaces[1], address='10.1.1.2/24'),
+            IPAddress(assigned_object=None, address='10.1.1.3/24'),
+            IPAddress(assigned_object=interfaces[0], address='2001:db8::1/64'),
+            IPAddress(assigned_object=interfaces[1], address='2001:db8::2/64'),
+            IPAddress(assigned_object=None, address='2001:db8::3/64'),
         )
         IPAddress.objects.bulk_create(addresses)
 
         vdcs[0].primary_ip4 = addresses[0]
+        vdcs[0].primary_ip6 = addresses[3]
         vdcs[0].save()
         vdcs[1].primary_ip4 = addresses[1]
+        vdcs[1].primary_ip6 = addresses[4]
         vdcs[1].save()
 
     def test_device(self):
@@ -4738,3 +4744,17 @@ def test_has_primary_ip(self):
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
         params = {'has_primary_ip': False}
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 4)
+
+    def test_primary_ip4(self):
+        addresses = IPAddress.objects.filter(address__family=4)
+        params = {'primary_ip4_id': [addresses[0].pk, addresses[1].pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+        params = {'primary_ip4_id': [addresses[2].pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 0)
+
+    def test_primary_ip6(self):
+        addresses = IPAddress.objects.filter(address__family=6)
+        params = {'primary_ip6_id': [addresses[0].pk, addresses[1].pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+        params = {'primary_ip6_id': [addresses[2].pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 0)

netbox/extras/tests/test_views.py

@@ -457,7 +457,7 @@ def setUpTestData(cls):
             'platforms': [],
             'tenant_groups': [],
             'tenants': [],
-            'device_types': [devicetype.id,],
+            'device_types': [devicetype.id],
             'tags': [],
             'data': '{"foo": 123}',
         }

netbox/ipam/filtersets.py

@@ -1227,3 +1227,19 @@ def filter_region(self, queryset, name, value):
             )
         )
         return qs
+
+
+class PrimaryIPFilterSet(django_filters.FilterSet):
+    """
+    An inheritable FilterSet for models which support primary IP assignment.
+    """
+    primary_ip4_id = django_filters.ModelMultipleChoiceFilter(
+        field_name='primary_ip4',
+        queryset=IPAddress.objects.all(),
+        label=_('Primary IPv4 (ID)'),
+    )
+    primary_ip6_id = django_filters.ModelMultipleChoiceFilter(
+        field_name='primary_ip6',
+        queryset=IPAddress.objects.all(),
+        label=_('Primary IPv6 (ID)'),
+    )

netbox/virtualization/filtersets.py

@@ -6,6 +6,7 @@
 from dcim.models import Device, DeviceRole, Platform, Region, Site, SiteGroup
 from extras.filtersets import LocalConfigContextFilterSet
 from extras.models import ConfigTemplate
+from ipam.filtersets import PrimaryIPFilterSet
 from netbox.filtersets import OrganizationalModelFilterSet, NetBoxModelFilterSet
 from tenancy.filtersets import TenancyFilterSet, ContactModelFilterSet
 from utilities.filters import MultiValueCharFilter, MultiValueMACAddressFilter, TreeNodeMultipleChoiceFilter
@@ -114,7 +115,8 @@ class VirtualMachineFilterSet(
     NetBoxModelFilterSet,
     TenancyFilterSet,
     ContactModelFilterSet,
-    LocalConfigContextFilterSet
+    LocalConfigContextFilterSet,
+    PrimaryIPFilterSet,
 ):
     status = django_filters.MultipleChoiceFilter(
         choices=VirtualMachineStatusChoices,

netbox/virtualization/tests/test_filtersets.py

@@ -291,10 +291,14 @@ def setUpTestData(cls):
         ipaddresses = (
             IPAddress(address='192.0.2.1/24', assigned_object=interfaces[0]),
             IPAddress(address='192.0.2.2/24', assigned_object=interfaces[1]),
+            IPAddress(address='192.0.2.3/24', assigned_object=None),
+            IPAddress(address='2001:db8::1/64', assigned_object=interfaces[0]),
+            IPAddress(address='2001:db8::2/64', assigned_object=interfaces[1]),
+            IPAddress(address='2001:db8::3/64', assigned_object=None),
         )
         IPAddress.objects.bulk_create(ipaddresses)
-        VirtualMachine.objects.filter(pk=vms[0].pk).update(primary_ip4=ipaddresses[0])
-        VirtualMachine.objects.filter(pk=vms[1].pk).update(primary_ip4=ipaddresses[1])
+        VirtualMachine.objects.filter(pk=vms[0].pk).update(primary_ip4=ipaddresses[0], primary_ip6=ipaddresses[3])
+        VirtualMachine.objects.filter(pk=vms[1].pk).update(primary_ip4=ipaddresses[1], primary_ip6=ipaddresses[4])
 
     def test_name(self):
         params = {'name': ['Virtual Machine 1', 'Virtual Machine 2']}
@@ -412,6 +416,20 @@ def test_tenant_group(self):
         params = {'tenant_group': [tenant_groups[0].slug, tenant_groups[1].slug]}
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
 
+    def test_primary_ip4(self):
+        addresses = IPAddress.objects.filter(address__family=4)
+        params = {'primary_ip4_id': [addresses[0].pk, addresses[1].pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+        params = {'primary_ip4_id': [addresses[2].pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 0)
+
+    def test_primary_ip6(self):
+        addresses = IPAddress.objects.filter(address__family=6)
+        params = {'primary_ip6_id': [addresses[0].pk, addresses[1].pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+        params = {'primary_ip6_id': [addresses[2].pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 0)
+
 
 class VMInterfaceTestCase(TestCase, ChangeLoggedFilterSetTests):
     queryset = VMInterface.objects.all()