From 73f4d4950eb504e3ca2b33b40ed73e2d2b8ed799 Mon Sep 17 00:00:00 2001
From: Jacob Foster <jacob@britecore.com>
Date: Wed, 19 Jul 2017 14:21:06 -0500
Subject: [PATCH] Add tests for optimized queries

---
 graphene_django/tests/test_optimization.py | 126 +++++++++++++++++++++
 1 file changed, 126 insertions(+)
 create mode 100644 graphene_django/tests/test_optimization.py

diff --git a/graphene_django/tests/test_optimization.py b/graphene_django/tests/test_optimization.py
new file mode 100644
index 0000000..05d898b
--- /dev/null
+++ b/graphene_django/tests/test_optimization.py
@@ -0,0 +1,126 @@
+from datetime import date
+from django.db import connection
+from django.test import TestCase
+from django.test.utils import CaptureQueriesContext
+import graphene
+import pytest
+
+from .. import registry
+from ..fields import DjangoConnectionField
+from ..optimization import optimize_queryset
+from ..types import DjangoObjectType
+from .models import (
+    Article as ArticleModel,
+    Reporter as ReporterModel,
+    Pet as PetModel
+)
+
+pytestmark = pytest.mark.django_db
+
+registry.reset_global_registry()
+
+
+class Article(DjangoObjectType):
+    class Meta:
+        model = ArticleModel
+        interfaces = (graphene.relay.Node,)
+
+
+class Reporter(DjangoObjectType):
+    class Meta:
+        model = ReporterModel
+
+
+class Pet(DjangoObjectType):
+    class Meta:
+        model = PetModel
+
+
+class RootQuery(graphene.ObjectType):
+    article = graphene.Field(Article, id=graphene.ID())
+    articles = DjangoConnectionField(Article)
+
+    def resolve_article(self, args, context, info):
+        qs = ArticleModel.objects
+        qs = optimize_queryset(ArticleModel, qs, info.field_asts[0])
+        return qs.get(**args)
+
+
+schema = graphene.Schema(query=RootQuery)
+
+
+class TestOptimization(TestCase):
+
+    @classmethod
+    def setUpTestData(cls):
+        cls.reporter = ReporterModel.objects.create(
+            first_name='Clark', last_name='Kent',
+            email='ckent@dailyplanet.com', a_choice='this'
+        )
+        cls.editor = ReporterModel.objects.create(
+            first_name='Perry', last_name='White',
+            email='pwhite@dailyplanet.com', a_choice='this'
+        )
+        cls.article = ArticleModel.objects.create(
+            headline='Superman Saves the Day',
+            pub_date=date.today(),
+            reporter=cls.reporter,
+            editor=cls.editor
+        )
+        cls.other_article = ArticleModel.objects.create(
+            headline='Lex Luthor is SO Rich',
+            pub_date=date.today(),
+            reporter=cls.reporter,
+            editor=cls.editor
+        )
+        cls.editor.pets.add(cls.reporter)
+
+    def test_select_related(self):
+        query = """query GetArticle($articleId: ID!){
+          article(id: $articleId) {
+              headline
+              reporter {
+                email
+              }
+              editor {
+                email
+              }
+          }
+        }"""
+
+        variables = {'articleId': str(self.article.id)}
+
+        with CaptureQueriesContext(connection) as query_context:
+            results = schema.execute(query, variable_values=variables)
+
+        returned_article = results.data['article']
+        assert returned_article['headline'] == self.article.headline
+        assert returned_article['reporter']['email'] == self.reporter.email
+        assert returned_article['editor']['email'] == self.editor.email
+
+        self.assertEqual(len(query_context.captured_queries), 1)
+
+    def test_prefetch_related(self):
+        query = """query {
+          articles {
+            edges {
+              node {
+                headline
+                editor {
+                  email
+                  pets {
+                    email
+                  }
+                }
+              }
+            }
+          }
+        }"""
+
+        with CaptureQueriesContext(connection) as query_context:
+            results = schema.execute(query)
+
+        returned_articles = results.data['articles']['edges']
+        assert len(returned_articles) == 2
+
+        self.assertEqual(len(query_context.captured_queries), 4)