From 2b75e05ea7d2b6ff0b50d7c991bcec593107326f Mon Sep 17 00:00:00 2001
From: Dan Sheppard <dans@dansheps.com>
Date: Sat, 22 Feb 2020 08:16:29 -0600
Subject: [PATCH] Fixes: #4230 - Fixes filtering by position on elevation
 endpoint

* Add tests for rack elevation filtering
* Add q variable to serializers for RackElevationDetailFilterSerializer
* Add code to allow filtering of position on the rack elevation
---
 netbox/dcim/api/serializers.py |  4 ++++
 netbox/dcim/api/views.py       |  5 +++++
 netbox/dcim/tests/test_api.py  | 22 ++++++++++++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/netbox/dcim/api/serializers.py b/netbox/dcim/api/serializers.py
index 34d9a7890..34f65a587 100644
--- a/netbox/dcim/api/serializers.py
+++ b/netbox/dcim/api/serializers.py
@@ -172,6 +172,10 @@ class RackReservationSerializer(ValidatedModelSerializer):
 
 
 class RackElevationDetailFilterSerializer(serializers.Serializer):
+    q = serializers.CharField(
+        required=False,
+        default=None
+    )
     face = serializers.ChoiceField(
         choices=DeviceFaceChoices,
         default=DeviceFaceChoices.FACE_FRONT
diff --git a/netbox/dcim/api/views.py b/netbox/dcim/api/views.py
index f8297fe46..c4a159126 100644
--- a/netbox/dcim/api/views.py
+++ b/netbox/dcim/api/views.py
@@ -237,6 +237,11 @@ class RackViewSet(CustomFieldModelViewSet):
                 expand_devices=data['expand_devices']
             )
 
+            # Enable filtering rack units by ID
+            q = data['q']
+            if q:
+                elevation = [u for u in elevation if q in str(u['id']) or q in str(u['name'])]
+
             page = self.paginate_queryset(elevation)
             if page is not None:
                 rack_units = serializers.RackUnitSerializer(page, many=True, context={'request': request})
diff --git a/netbox/dcim/tests/test_api.py b/netbox/dcim/tests/test_api.py
index a3a072bc9..ff86d6ccc 100644
--- a/netbox/dcim/tests/test_api.py
+++ b/netbox/dcim/tests/test_api.py
@@ -596,6 +596,28 @@ class RackTest(APITestCase):
 
         self.assertEqual(response.data['count'], 42)
 
+    def test_get_elevation_rack_units(self):
+
+        url = '{}?q=3'.format(reverse('dcim-api:rack-elevation', kwargs={'pk': self.rack1.pk}))
+        response = self.client.get(url, **self.header)
+
+        self.assertEqual(response.data['count'], 13)
+
+        url = '{}?q=U3'.format(reverse('dcim-api:rack-elevation', kwargs={'pk': self.rack1.pk}))
+        response = self.client.get(url, **self.header)
+
+        self.assertEqual(response.data['count'], 11)
+
+        url = '{}?q=10'.format(reverse('dcim-api:rack-elevation', kwargs={'pk': self.rack1.pk}))
+        response = self.client.get(url, **self.header)
+
+        self.assertEqual(response.data['count'], 1)
+
+        url = '{}?q=U20'.format(reverse('dcim-api:rack-elevation', kwargs={'pk': self.rack1.pk}))
+        response = self.client.get(url, **self.header)
+
+        self.assertEqual(response.data['count'], 1)
+
     def test_get_rack_elevation(self):
 
         url = reverse('dcim-api:rack-elevation', kwargs={'pk': self.rack1.pk})