Add device & vm to `FHRPGroupAssignmentFilterSet` (#9314)

* Add device & vm to `FHRPGroupAssignmentFilterSet`

* Apply suggestions from code review

* Update netbox/ipam/tests/test_filtersets.py

* Update netbox/ipam/filtersets.py

Co-authored-by: Jeremy Stretch <jstretch@ns1.com>

netbox/ipam/filtersets.py

@@ -681,11 +681,53 @@ class FHRPGroupAssignmentFilterSet(ChangeLoggedModelFilterSet):
         queryset=FHRPGroup.objects.all(),
         label='Group (ID)',
     )
+    device = MultiValueCharFilter(
+        method='filter_device',
+        field_name='name',
+        label='Device (name)',
+    )
+    device_id = MultiValueNumberFilter(
+        method='filter_device',
+        field_name='pk',
+        label='Device (ID)',
+    )
+    virtual_machine = MultiValueCharFilter(
+        method='filter_virtual_machine',
+        field_name='name',
+        label='Virtual machine (name)',
+    )
+    virtual_machine_id = MultiValueNumberFilter(
+        method='filter_virtual_machine',
+        field_name='pk',
+        label='Virtual machine (ID)',
+    )
 
     class Meta:
         model = FHRPGroupAssignment
         fields = ['id', 'group_id', 'interface_type', 'interface_id', 'priority']
 
+    def filter_device(self, queryset, name, value):
+        devices = Device.objects.filter(**{f'{name}__in': value})
+        if not devices.exists():
+            return queryset.none()
+        interface_ids = []
+        for device in devices:
+            interface_ids.extend(device.vc_interfaces().values_list('id', flat=True))
+        return queryset.filter(
+            Q(interface_type=ContentType.objects.get_for_model(Interface), interface_id__in=interface_ids)
+        )
+
+    def filter_virtual_machine(self, queryset, name, value):
+        virtual_machines = VirtualMachine.objects.filter(**{f'{name}__in': value})
+        if not virtual_machines.exists():
+            return queryset.none()
+        interface_ids = []
+        for vm in virtual_machines:
+            interface_ids.extend(vm.interfaces.values_list('id', flat=True))
+        return queryset.filter(
+            Q(interface_type=ContentType.objects.get_for_model(VMInterface), interface_id__in=interface_ids)
+        )
+
 
 class VLANGroupFilterSet(OrganizationalModelFilterSet):
     scope_type = ContentTypeFilter()

netbox/ipam/tests/test_filtersets.py

@@ -1024,6 +1024,20 @@ class FHRPGroupAssignmentTestCase(TestCase, ChangeLoggedFilterSetTests):
         params = {'priority': [10, 20]}
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 4)
 
+    def test_device(self):
+        device = Device.objects.first()
+        params = {'device': [device.name]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 3)
+        params = {'device_id': [device.pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 3)
+
+    def test_virtual_machine(self):
+        vm = VirtualMachine.objects.first()
+        params = {'virtual_machine': [vm.name]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 3)
+        params = {'virtual_machine_id': [vm.pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 3)
+
 
 class VLANGroupTestCase(TestCase, ChangeLoggedFilterSetTests):
     queryset = VLANGroup.objects.all()