From bddc24ef5641c346920b3efb4ea696ed0cf90908 Mon Sep 17 00:00:00 2001
From: rmanyari <rmanyari10@gmail.com>
Date: Wed, 1 Mar 2023 15:35:40 -0700
Subject: [PATCH] Add family field support to Prefix and Aggregate, fix tests

---
 netbox/ipam/graphql/types.py    | 44 +++++++++++++++++++--------------
 netbox/utilities/testing/api.py |  4 +++
 2 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/netbox/ipam/graphql/types.py b/netbox/ipam/graphql/types.py
index cdc2e667a..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,30 +86,14 @@ class FHRPGroupAssignmentType(BaseObjectType):
         filterset_class = filtersets.FHRPGroupAssignmentFilterSet
 
 
-class IPAddressFamilyType(graphene.ObjectType):
-
-    value = graphene.Int()
-    label = graphene.String()
-
-    def __init__(self, value):
-        self.value = value
-        self.label = 'IPv4' if value == 4 else 'IPv6'
-
-
-class IPAddressType(NetBoxObjectType):
+class IPAddressType(NetBoxObjectType, BaseIPAddressFamilyType):
     assigned_object = graphene.Field('ipam.graphql.gfk_mixins.IPAddressAssignmentType')
-    family = graphene.Field(IPAddressFamilyType)
 
     class Meta:
         model = models.IPAddress
         exclude = ('assigned_object_type', 'assigned_object_id')
         filterset_class = filtersets.IPAddressFilterSet
 
-    def resolve_family(self, _):
-        # Note that self, is an instance of models.IPAddress
-        # thus resolves to the address family value.
-        return IPAddressFamilyType(self.family)
-
     def resolve_role(self, info):
         return self.role or None
 
@@ -103,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'