From ab303db3dde9be39385fb13028b0896aee2e6d6a Mon Sep 17 00:00:00 2001
From: rmanyari <rmanyari@users.noreply.github.com>
Date: Fri, 10 Mar 2023 12:48:45 -0700
Subject: [PATCH] Closes #11851: Add family field to IPAddress queries in
 GraphQL (#11870)

* Closes #11851: Add family field to IPAddress queries in GraphQL

* Add family field support to Prefix and Aggregate, fix tests
---
 netbox/ipam/graphql/types.py    | 28 +++++++++++++++++++++++++---
 netbox/utilities/testing/api.py |  4 ++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/netbox/ipam/graphql/types.py b/netbox/ipam/graphql/types.py
index b8f6221bc..0da61ea8a 100644
--- a/netbox/ipam/graphql/types.py
+++ b/netbox/ipam/graphql/types.py
@@ -27,6 +27,28 @@ __all__ = (
 )
 
 
+class IPAddressFamilyType(graphene.ObjectType):
+
+    value = graphene.Int()
+    label = graphene.String()
+
+    def __init__(self, value):
+        self.value = value
+        self.label = f'IPv{value}'
+
+
+class BaseIPAddressFamilyType:
+    '''
+    Base type for models that need to expose their IPAddress family type.
+    '''
+    family = graphene.Field(IPAddressFamilyType)
+
+    def resolve_family(self, _):
+        # Note that self, is an instance of models.IPAddress
+        # thus resolves to the address family value.
+        return IPAddressFamilyType(self.family)
+
+
 class ASNType(NetBoxObjectType):
     asn = graphene.Field(BigInt)
 
@@ -36,7 +58,7 @@ class ASNType(NetBoxObjectType):
         filterset_class = filtersets.ASNFilterSet
 
 
-class AggregateType(NetBoxObjectType):
+class AggregateType(NetBoxObjectType, BaseIPAddressFamilyType):
 
     class Meta:
         model = models.Aggregate
@@ -64,7 +86,7 @@ class FHRPGroupAssignmentType(BaseObjectType):
         filterset_class = filtersets.FHRPGroupAssignmentFilterSet
 
 
-class IPAddressType(NetBoxObjectType):
+class IPAddressType(NetBoxObjectType, BaseIPAddressFamilyType):
     assigned_object = graphene.Field('ipam.graphql.gfk_mixins.IPAddressAssignmentType')
 
     class Meta:
@@ -87,7 +109,7 @@ class IPRangeType(NetBoxObjectType):
         return self.role or None
 
 
-class PrefixType(NetBoxObjectType):
+class PrefixType(NetBoxObjectType, BaseIPAddressFamilyType):
 
     class Meta:
         model = models.Prefix
diff --git a/netbox/utilities/testing/api.py b/netbox/utilities/testing/api.py
index c0836b71d..7f24c86b8 100644
--- a/netbox/utilities/testing/api.py
+++ b/netbox/utilities/testing/api.py
@@ -17,6 +17,8 @@ from utilities.api import get_graphql_type_for_model
 from .base import ModelTestCase
 from .utils import disable_warnings
 
+from ipam.graphql.types import IPAddressFamilyType
+
 
 __all__ = (
     'APITestCase',
@@ -460,6 +462,8 @@ class APIViewTestCases:
                     # TODO: Come up with something more elegant
                     # Temporary hack to support automated testing of reverse generic relations
                     fields_string += f'{field_name} {{ id }}\n'
+                elif inspect.isclass(field.type) and issubclass(field.type, IPAddressFamilyType):
+                    fields_string += f'{field_name} {{ value, label }}\n'
                 else:
                     fields_string += f'{field_name}\n'