Add family field support to Prefix and Aggregate, fix tests

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

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'