From cc721efe97253e8f0e9345fd335411bad1a6de09 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Tue, 21 Apr 2020 14:12:49 -0400
Subject: [PATCH] Fixes #3356: Correct Swagger schema specification for the
 available prefixes/IPs API endpoints

---
 docs/release-notes/version-2.8.md |  1 +
 netbox/ipam/api/serializers.py    |  8 ++++++++
 netbox/ipam/api/views.py          | 13 +++++++++++++
 3 files changed, 22 insertions(+)

diff --git a/docs/release-notes/version-2.8.md b/docs/release-notes/version-2.8.md
index e29ee1aa1..0f4e0ba48 100644
--- a/docs/release-notes/version-2.8.md
+++ b/docs/release-notes/version-2.8.md
@@ -9,6 +9,7 @@
 ### Bug Fixes
 
 * [#2994](https://github.com/netbox-community/netbox/issues/2994) - Prevent modifying termination points of existing cable to ensure end-to-end path integrity
+* [#3356](https://github.com/netbox-community/netbox/issues/3356) - Correct Swagger schema specification for the available prefixes/IPs API endpoints
 * [#4361](https://github.com/netbox-community/netbox/issues/4361) - Fix Type of `connection_state` in Swagger schema
 * [#4388](https://github.com/netbox-community/netbox/issues/4388) - Fix detection of connected endpoints when connecting rear ports
 * [#4489](https://github.com/netbox-community/netbox/issues/4489) - Fix display of parent/child role on device type view
diff --git a/netbox/ipam/api/serializers.py b/netbox/ipam/api/serializers.py
index 4e596631d..f5de2f509 100644
--- a/netbox/ipam/api/serializers.py
+++ b/netbox/ipam/api/serializers.py
@@ -183,6 +183,10 @@ class AvailablePrefixSerializer(serializers.Serializer):
     """
     Representation of a prefix which does not exist in the database.
     """
+    family = serializers.IntegerField(read_only=True)
+    prefix = serializers.CharField(read_only=True)
+    vrf = NestedVRFSerializer(read_only=True)
+
     def to_representation(self, instance):
         if self.context.get('vrf'):
             vrf = NestedVRFSerializer(self.context['vrf'], context={'request': self.context['request']}).data
@@ -246,6 +250,10 @@ class AvailableIPSerializer(serializers.Serializer):
     """
     Representation of an IP address which does not exist in the database.
     """
+    family = serializers.IntegerField(read_only=True)
+    address = serializers.CharField(read_only=True)
+    vrf = NestedVRFSerializer(read_only=True)
+
     def to_representation(self, instance):
         if self.context.get('vrf'):
             vrf = NestedVRFSerializer(self.context['vrf'], context={'request': self.context['request']}).data
diff --git a/netbox/ipam/api/views.py b/netbox/ipam/api/views.py
index f24c71b17..bf430f633 100644
--- a/netbox/ipam/api/views.py
+++ b/netbox/ipam/api/views.py
@@ -2,6 +2,7 @@ from django.conf import settings
 from django.db.models import Count
 from django.shortcuts import get_object_or_404
 from django_pglocks import advisory_lock
+from drf_yasg.utils import swagger_auto_schema
 from rest_framework import status
 from rest_framework.decorators import action
 from rest_framework.exceptions import PermissionDenied
@@ -73,6 +74,12 @@ class PrefixViewSet(CustomFieldModelViewSet):
     serializer_class = serializers.PrefixSerializer
     filterset_class = filters.PrefixFilterSet
 
+    @swagger_auto_schema(
+        methods=['get', 'post'],
+        responses={
+            200: serializers.AvailablePrefixSerializer(many=True),
+        }
+    )
     @action(detail=True, url_path='available-prefixes', methods=['get', 'post'])
     @advisory_lock(ADVISORY_LOCK_KEYS['available-prefixes'])
     def available_prefixes(self, request, pk=None):
@@ -151,6 +158,12 @@ class PrefixViewSet(CustomFieldModelViewSet):
 
             return Response(serializer.data)
 
+    @swagger_auto_schema(
+        methods=['get', 'post'],
+        responses={
+            200: serializers.AvailableIPSerializer(many=True),
+        }
+    )
     @action(detail=True, url_path='available-ips', methods=['get', 'post'])
     @advisory_lock(ADVISORY_LOCK_KEYS['available-ips'])
     def available_ips(self, request, pk=None):