Просмотр исходного кода

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
rmanyari 2 лет назад
Родитель
Сommit
ab303db3dd
2 измененных файлов с 29 добавлено и 3 удалено
  1. 25 3
      netbox/ipam/graphql/types.py
  2. 4 0
      netbox/utilities/testing/api.py

+ 25 - 3
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

+ 4 - 0
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'