From 10bb8fa10af10e976ff55e15b0cbcf87a4002b33 Mon Sep 17 00:00:00 2001
From: Arthur Hanson <worldnomad@gmail.com>
Date: Thu, 6 Oct 2022 13:50:53 -0700
Subject: [PATCH] 9478 link peers to GraphQL (#10574)

* 9468 add link_peer to GraphQL

* 9478 add class_type

* 9478 fix tests

* 9478 fix tests

* 9478 fix tests
---
 netbox/dcim/graphql/gfk_mixins.py | 59 +++++++++++++++++++++++++++++++
 netbox/dcim/graphql/mixins.py     |  7 ++++
 netbox/netbox/graphql/types.py    | 18 +++++++---
 netbox/utilities/testing/api.py   |  6 +++-
 4 files changed, 84 insertions(+), 6 deletions(-)
 create mode 100644 netbox/dcim/graphql/gfk_mixins.py

diff --git a/netbox/dcim/graphql/gfk_mixins.py b/netbox/dcim/graphql/gfk_mixins.py
new file mode 100644
index 000000000..d6be138bc
--- /dev/null
+++ b/netbox/dcim/graphql/gfk_mixins.py
@@ -0,0 +1,59 @@
+import graphene
+from circuits.graphql.types import CircuitTerminationType
+from circuits.models import CircuitTermination
+from dcim.graphql.types import (
+    ConsolePortType,
+    ConsoleServerPortType,
+    FrontPortType,
+    InterfaceType,
+    PowerFeedType,
+    PowerOutletType,
+    PowerPortType,
+    RearPortType,
+)
+from dcim.models import (
+    ConsolePort,
+    ConsoleServerPort,
+    FrontPort,
+    Interface,
+    PowerFeed,
+    PowerOutlet,
+    PowerPort,
+    RearPort,
+)
+
+
+class LinkPeerType(graphene.Union):
+    class Meta:
+        types = (
+            CircuitTerminationType,
+            ConsolePortType,
+            ConsoleServerPortType,
+            FrontPortType,
+            InterfaceType,
+            PowerFeedType,
+            PowerOutletType,
+            PowerPortType,
+            RearPortType,
+        )
+
+    @classmethod
+    def resolve_type(cls, instance, info):
+        if type(instance) == CircuitTermination:
+            return CircuitTerminationType
+        if type(instance) == ConsolePortType:
+            return ConsolePortType
+        if type(instance) == ConsoleServerPort:
+            return ConsoleServerPortType
+        if type(instance) == FrontPort:
+            return FrontPortType
+        if type(instance) == Interface:
+            return InterfaceType
+        if type(instance) == PowerFeed:
+            return PowerFeedType
+        if type(instance) == PowerOutlet:
+            return PowerOutletType
+        if type(instance) == PowerPort:
+            return PowerPortType
+        if type(instance) == RearPort:
+            return RearPortType
diff --git a/netbox/dcim/graphql/mixins.py b/netbox/dcim/graphql/mixins.py
index d8488aa5f..133d6259f 100644
--- a/netbox/dcim/graphql/mixins.py
+++ b/netbox/dcim/graphql/mixins.py
@@ -1,5 +1,12 @@
+import graphene
+
+
 class CabledObjectMixin:
+    link_peers = graphene.List('dcim.graphql.gfk_mixins.LinkPeerType')
 
     def resolve_cable_end(self, info):
         # Handle empty values
         return self.cable_end or None
+
+    def resolve_link_peers(self, info):
+        return self.link_peers
diff --git a/netbox/netbox/graphql/types.py b/netbox/netbox/graphql/types.py
index 41eff6d46..10847742b 100644
--- a/netbox/netbox/graphql/types.py
+++ b/netbox/netbox/graphql/types.py
@@ -1,10 +1,14 @@
 import graphene
 
 from django.contrib.contenttypes.models import ContentType
+from extras.graphql.mixins import (
+    ChangelogMixin,
+    CustomFieldsMixin,
+    JournalEntriesMixin,
+    TagsMixin,
+)
 from graphene_django import DjangoObjectType
 
-from extras.graphql.mixins import ChangelogMixin, CustomFieldsMixin, JournalEntriesMixin, TagsMixin
-
 __all__ = (
     'BaseObjectType',
     'ObjectType',
@@ -22,9 +26,7 @@ class BaseObjectType(DjangoObjectType):
     Base GraphQL object type for all NetBox objects. Restricts the model queryset to enforce object permissions.
     """
     display = graphene.String()
-
-    def resolve_display(parent, info, **kwargs):
-        return str(parent)
+    class_type = graphene.String()
 
     class Meta:
         abstract = True
@@ -34,6 +36,12 @@ class BaseObjectType(DjangoObjectType):
         # Enforce object permissions on the queryset
         return queryset.restrict(info.context.user, 'view')
 
+    def resolve_display(parent, info, **kwargs):
+        return str(parent)
+
+    def resolve_class_type(parent, info, **kwargs):
+        return parent.__class__.__name__
+
 
 class ObjectType(
     ChangelogMixin,
diff --git a/netbox/utilities/testing/api.py b/netbox/utilities/testing/api.py
index f26e5fffc..8815ede1f 100644
--- a/netbox/utilities/testing/api.py
+++ b/netbox/utilities/testing/api.py
@@ -1,3 +1,4 @@
+import inspect
 import json
 
 from django.conf import settings
@@ -5,7 +6,7 @@ from django.contrib.auth.models import User
 from django.contrib.contenttypes.models import ContentType
 from django.urls import reverse
 from django.test import override_settings
-from graphene.types import Dynamic as GQLDynamic, List as GQLList
+from graphene.types import Dynamic as GQLDynamic, List as GQLList, Union as GQLUnion
 from rest_framework import status
 from rest_framework.test import APIClient
 
@@ -449,6 +450,9 @@ class APIViewTestCases:
                 if type(field) is GQLDynamic:
                     # Dynamic fields must specify a subselection
                     fields_string += f'{field_name} {{ id }}\n'
+                elif type(field.type) is GQLList and inspect.isclass(field.type.of_type) and issubclass(field.type.of_type, GQLUnion):
+                    # Union types dont' have an id or consistent values
+                    continue
                 elif type(field.type) is GQLList and field_name != 'choices':
                     # TODO: Come up with something more elegant
                     # Temporary hack to support automated testing of reverse generic relations