From 881b18f6d0361ff2a882d74933420ffa3b866dba Mon Sep 17 00:00:00 2001
From: jeremystretch <jstretch@ns1.com>
Date: Fri, 25 Jun 2021 14:19:20 -0400
Subject: [PATCH] Add GraphQL for tenancy

---
 netbox/netbox/graphql/schema.py    |  2 ++
 netbox/tenancy/graphql/__init__.py |  0
 netbox/tenancy/graphql/schema.py   | 12 ++++++++++++
 netbox/tenancy/graphql/types.py    | 23 +++++++++++++++++++++++
 netbox/tenancy/tests/test_api.py   |  4 ++--
 5 files changed, 39 insertions(+), 2 deletions(-)
 create mode 100644 netbox/tenancy/graphql/__init__.py
 create mode 100644 netbox/tenancy/graphql/schema.py
 create mode 100644 netbox/tenancy/graphql/types.py

diff --git a/netbox/netbox/graphql/schema.py b/netbox/netbox/graphql/schema.py
index 0899bd564..dcb1107e1 100644
--- a/netbox/netbox/graphql/schema.py
+++ b/netbox/netbox/graphql/schema.py
@@ -3,12 +3,14 @@ import graphene
 from circuits.graphql.schema import CircuitsQuery
 from extras.graphql.schema import ExtrasQuery
 from ipam.graphql.schema import IPAMQuery
+from tenancy.graphql.schema import TenancyQuery
 
 
 class Query(
     CircuitsQuery,
     ExtrasQuery,
     IPAMQuery,
+    TenancyQuery,
     graphene.ObjectType
 ):
     pass
diff --git a/netbox/tenancy/graphql/__init__.py b/netbox/tenancy/graphql/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/netbox/tenancy/graphql/schema.py b/netbox/tenancy/graphql/schema.py
new file mode 100644
index 000000000..5a2f7438b
--- /dev/null
+++ b/netbox/tenancy/graphql/schema.py
@@ -0,0 +1,12 @@
+import graphene
+
+from netbox.graphql.fields import ObjectField, ObjectListField
+from .types import *
+
+
+class TenancyQuery(graphene.ObjectType):
+    tenant = ObjectField(TenantType)
+    tenants = ObjectListField(TenantType)
+
+    tenant_group = ObjectField(TenantGroupType)
+    tenant_groups = ObjectListField(TenantGroupType)
diff --git a/netbox/tenancy/graphql/types.py b/netbox/tenancy/graphql/types.py
new file mode 100644
index 000000000..ca8e839d8
--- /dev/null
+++ b/netbox/tenancy/graphql/types.py
@@ -0,0 +1,23 @@
+from tenancy import filtersets, models
+from netbox.graphql.types import ObjectType, TaggedObjectType
+
+__all__ = (
+    'TenantType',
+    'TenantGroupType',
+)
+
+
+class TenantType(ObjectType):
+
+    class Meta:
+        model = models.Tenant
+        fields = '__all__'
+        filterset_class = filtersets.TenantFilterSet
+
+
+class TenantGroupType(TaggedObjectType):
+
+    class Meta:
+        model = models.TenantGroup
+        fields = '__all__'
+        filterset_class = filtersets.TenantGroupFilterSet
diff --git a/netbox/tenancy/tests/test_api.py b/netbox/tenancy/tests/test_api.py
index 5a3c2c1b0..1777f204a 100644
--- a/netbox/tenancy/tests/test_api.py
+++ b/netbox/tenancy/tests/test_api.py
@@ -14,7 +14,7 @@ class AppTest(APITestCase):
         self.assertEqual(response.status_code, 200)
 
 
-class TenantGroupTest(APIViewTestCases.APIViewTestCase):
+class TenantGroupTest(APIViewTestCases.GraphQLTestCase, APIViewTestCases.APIViewTestCase):
     model = TenantGroup
     brief_fields = ['_depth', 'display', 'id', 'name', 'slug', 'tenant_count', 'url']
     bulk_update_data = {
@@ -52,7 +52,7 @@ class TenantGroupTest(APIViewTestCases.APIViewTestCase):
         ]
 
 
-class TenantTest(APIViewTestCases.APIViewTestCase):
+class TenantTest(APIViewTestCases.GraphQLTestCase, APIViewTestCases.APIViewTestCase):
     model = Tenant
     brief_fields = ['display', 'id', 'name', 'slug', 'url']
     bulk_update_data = {