From 0b38d40b48cd4e0376979c971c75f05e19d3306f Mon Sep 17 00:00:00 2001
From: Arthur <worldnomad@gmail.com>
Date: Thu, 6 Oct 2022 10:07:09 -0700
Subject: [PATCH] 9817 ipan graphql gfk

---
 netbox/ipam/graphql/gfk_mixins.py | 62 ++++++++++++++++++++++++++++---
 netbox/ipam/graphql/types.py      |  6 ++-
 2 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/netbox/ipam/graphql/gfk_mixins.py b/netbox/ipam/graphql/gfk_mixins.py
index c9ca617d2..31742c4a4 100644
--- a/netbox/ipam/graphql/gfk_mixins.py
+++ b/netbox/ipam/graphql/gfk_mixins.py
@@ -1,10 +1,17 @@
 import graphene
-from dcim.graphql.types import InterfaceType
-from dcim.models import Interface
+from dcim.graphql.types import (
+    InterfaceType,
+    LocationType,
+    RackType,
+    RegionType,
+    SiteGroupType,
+    SiteType,
+)
+from dcim.models import Interface, Location, Rack, Region, Site, SiteGroup
 from ipam.graphql.types import FHRPGroupType, VLANType
-from ipam.models import FHRPGroup, VLAN
-from virtualization.graphql.types import VMInterfaceType
-from virtualization.models import VMInterface
+from ipam.models import VLAN, FHRPGroup
+from virtualization.graphql.types import ClusterGroupType, ClusterType, VMInterfaceType
+from virtualization.models import Cluster, ClusterGroup, VMInterface
 
 
 class IPAddressAssignmentType(graphene.Union):
@@ -41,3 +48,48 @@ class L2VPNAssignmentType(graphene.Union):
             return VLANType
         if type(instance) == VMInterface:
             return VMInterfaceType
+
+
+class FHRPGroupInterfaceType(graphene.Union):
+    class Meta:
+        types = (
+            InterfaceType,
+            VMInterfaceType,
+        )
+
+    @classmethod
+    def resolve_type(cls, instance, info):
+        if type(instance) == Interface:
+            return InterfaceType
+        if type(instance) == VMInterface:
+            return VMInterfaceType
+
+
+class VLANGroupScopeType(graphene.Union):
+    class Meta:
+        types = (
+            ClusterType,
+            ClusterGroupType,
+            LocationType,
+            RackType,
+            RegionType,
+            SiteType,
+            SiteGroupType,
+        )
+
+    @classmethod
+    def resolve_type(cls, instance, info):
+        if type(instance) == Cluster:
+            return ClusterType
+        if type(instance) == ClusterGroup:
+            return ClusterGroupType
+        if type(instance) == Location:
+            return LocationType
+        if type(instance) == Rack:
+            return RackType
+        if type(instance) == Region:
+            return RegionType
+        if type(instance) == Site:
+            return SiteType
+        if type(instance) == SiteGroup:
+            return SiteGroupType
diff --git a/netbox/ipam/graphql/types.py b/netbox/ipam/graphql/types.py
index 4514c3acb..4937b6a9f 100644
--- a/netbox/ipam/graphql/types.py
+++ b/netbox/ipam/graphql/types.py
@@ -55,10 +55,11 @@ class FHRPGroupType(NetBoxObjectType):
 
 
 class FHRPGroupAssignmentType(BaseObjectType):
+    interface = graphene.Field('ipam.graphql.gfk_mixins.FHRPGroupInterfaceType')
 
     class Meta:
         model = models.FHRPGroupAssignment
-        fields = '__all__'
+        exclude = ('interface_type', 'interface_id')
         filterset_class = filtersets.FHRPGroupAssignmentFilterSet
 
 
@@ -142,10 +143,11 @@ class VLANType(NetBoxObjectType):
 
 
 class VLANGroupType(OrganizationalObjectType):
+    scope = graphene.Field('ipam.graphql.gfk_mixins.VLANGroupScopeType')
 
     class Meta:
         model = models.VLANGroup
-        fields = '__all__'
+        exclude = ('scope_type', 'scope_id')
         filterset_class = filtersets.VLANGroupFilterSet