From de17a651e6f976e8b7c16b49d4e78f6a6988b870 Mon Sep 17 00:00:00 2001
From: jeremystretch <jstretch@ns1.com>
Date: Tue, 30 Aug 2022 14:52:12 -0400
Subject: [PATCH] Closes #10043: Add support for 'limit' query parameter to
 available VLANs API endpoint

---
 docs/release-notes/version-3.3.md |  1 +
 netbox/ipam/api/views.py          | 29 ++++++++++++++++++-----------
 netbox/ipam/tests/test_api.py     | 23 ++++++++++++++++++-----
 3 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/docs/release-notes/version-3.3.md b/docs/release-notes/version-3.3.md
index bc091dd65..5298c1d26 100644
--- a/docs/release-notes/version-3.3.md
+++ b/docs/release-notes/version-3.3.md
@@ -4,6 +4,7 @@
 
 ### Enhancements
 
+* [#10043](https://github.com/netbox-community/netbox/issues/10043) - Add support for `limit` query parameter to available VLANs API endpoint
 * [#10060](https://github.com/netbox-community/netbox/issues/10060) - Add journal entries to global search
 
 ### Bug Fixes
diff --git a/netbox/ipam/api/views.py b/netbox/ipam/api/views.py
index 0407c6d39..9db3d7953 100644
--- a/netbox/ipam/api/views.py
+++ b/netbox/ipam/api/views.py
@@ -174,6 +174,21 @@ class L2VPNTerminationViewSet(NetBoxModelViewSet):
 # Views
 #
 
+def get_results_limit(request):
+    """
+    Return the lesser of the specified limit (if any) and the configured MAX_PAGE_SIZE.
+    """
+    config = get_config()
+    try:
+        limit = int(request.query_params.get('limit', config.PAGINATE_COUNT)) or config.MAX_PAGE_SIZE
+    except ValueError:
+        limit = config.PAGINATE_COUNT
+    if config.MAX_PAGE_SIZE:
+        limit = min(limit, config.MAX_PAGE_SIZE)
+
+    return limit
+
+
 class AvailablePrefixesView(ObjectValidationMixin, APIView):
     queryset = Prefix.objects.all()
 
@@ -265,16 +280,7 @@ class AvailableIPAddressesView(ObjectValidationMixin, APIView):
     @swagger_auto_schema(responses={200: serializers.AvailableIPSerializer(many=True)})
     def get(self, request, pk):
         parent = self.get_parent(request, pk)
-        config = get_config()
-        PAGINATE_COUNT = config.PAGINATE_COUNT
-        MAX_PAGE_SIZE = config.MAX_PAGE_SIZE
-
-        try:
-            limit = int(request.query_params.get('limit', PAGINATE_COUNT))
-        except ValueError:
-            limit = PAGINATE_COUNT
-        if MAX_PAGE_SIZE:
-            limit = min(limit, MAX_PAGE_SIZE)
+        limit = get_results_limit(request)
 
         # Calculate available IPs within the parent
         ip_list = []
@@ -357,8 +363,9 @@ class AvailableVLANsView(ObjectValidationMixin, APIView):
     @swagger_auto_schema(responses={200: serializers.AvailableVLANSerializer(many=True)})
     def get(self, request, pk):
         vlangroup = get_object_or_404(VLANGroup.objects.restrict(request.user), pk=pk)
-        available_vlans = vlangroup.get_available_vids()
+        limit = get_results_limit(request)
 
+        available_vlans = vlangroup.get_available_vids()[:limit]
         serializer = serializers.AvailableVLANSerializer(available_vlans, many=True, context={
             'request': request,
             'group': vlangroup,
diff --git a/netbox/ipam/tests/test_api.py b/netbox/ipam/tests/test_api.py
index 5dc708cd0..0fefb0162 100644
--- a/netbox/ipam/tests/test_api.py
+++ b/netbox/ipam/tests/test_api.py
@@ -699,9 +699,18 @@ class VLANGroupTest(APIViewTestCases.APIViewTestCase):
         """
         Test retrieval of all available VLANs within a group.
         """
-        self.add_permissions('ipam.view_vlangroup', 'ipam.view_vlan')
-        vlangroup = VLANGroup.objects.first()
+        MIN_VID = 100
+        MAX_VID = 199
 
+        self.add_permissions('ipam.view_vlangroup', 'ipam.view_vlan')
+        vlangroup = VLANGroup.objects.create(
+            name='VLAN Group X',
+            slug='vlan-group-x',
+            min_vid=MIN_VID,
+            max_vid=MAX_VID
+        )
+
+        # Create a set of VLANs within the group
         vlans = (
             VLAN(vid=10, name='VLAN 10', group=vlangroup),
             VLAN(vid=20, name='VLAN 20', group=vlangroup),
@@ -711,13 +720,17 @@ class VLANGroupTest(APIViewTestCases.APIViewTestCase):
 
         # Retrieve all available VLANs
         url = reverse('ipam-api:vlangroup-available-vlans', kwargs={'pk': vlangroup.pk})
-        response = self.client.get(url, **self.header)
-
-        self.assertEqual(len(response.data), 4094 - len(vlans))
+        response = self.client.get(f'{url}?limit=0', **self.header)
+        self.assertEqual(len(response.data), MAX_VID - MIN_VID + 1)
         available_vlans = {vlan['vid'] for vlan in response.data}
         for vlan in vlans:
             self.assertNotIn(vlan.vid, available_vlans)
 
+        # Retrieve a maximum number of available VLANs
+        url = reverse('ipam-api:vlangroup-available-vlans', kwargs={'pk': vlangroup.pk})
+        response = self.client.get(f'{url}?limit=10', **self.header)
+        self.assertEqual(len(response.data), 10)
+
     def test_create_single_available_vlan(self):
         """
         Test the creation of a single available VLAN.