bugfix: querytraversal should also work with mutations and subscriptions

andimarek · andimarek · commit 0e6edd285fa9 · 2017-08-30T14:59:10.000+02:00
diff --git a/src/main/java/graphql/analysis/QueryTraversal.java b/src/main/java/graphql/analysis/QueryTraversal.java
@@ -16,11 +16,15 @@
 import graphql.schema.GraphQLCompositeType;
 import graphql.schema.GraphQLFieldDefinition;
 import graphql.schema.GraphQLFieldsContainer;
+import graphql.schema.GraphQLObjectType;
 import graphql.schema.GraphQLSchema;
 
 import java.util.LinkedHashMap;
 import java.util.Map;
 
+import static graphql.Assert.assertNotNull;
+import static graphql.Assert.assertShouldNeverHappen;
+
 @Internal
 public class QueryTraversal {
 
@@ -47,11 +51,23 @@ public QueryTraversal(GraphQLSchema schema,
     }
 
     public void visitPostOrder(QueryVisitor visitor) {
-        visitImpl(visitor, operationDefinition.getSelectionSet(), schema.getQueryType(), null, false);
+        visitImpl(visitor, operationDefinition.getSelectionSet(), getRootType(), null, false);
     }
 
     public void visitPreOrder(QueryVisitor visitor) {
-        visitImpl(visitor, operationDefinition.getSelectionSet(), schema.getQueryType(), null, true);
+        visitImpl(visitor, operationDefinition.getSelectionSet(), getRootType(), null, true);
+    }
+
+    private GraphQLObjectType getRootType() {
+        if (operationDefinition.getOperation() == OperationDefinition.Operation.MUTATION) {
+            return assertNotNull(schema.getMutationType());
+        } else if (operationDefinition.getOperation() == OperationDefinition.Operation.QUERY) {
+            return assertNotNull(schema.getQueryType());
+        } else if (operationDefinition.getOperation() == OperationDefinition.Operation.SUBSCRIPTION) {
+            return assertNotNull(schema.getSubscriptionType());
+        } else {
+            return assertShouldNeverHappen();
+        }
     }
 
     public <T> T reducePostOrder(QueryReducer<T> queryReducer, T initialValue) {
diff --git a/src/test/groovy/graphql/analysis/QueryTraversalTest.groovy b/src/test/groovy/graphql/analysis/QueryTraversalTest.groovy
@@ -89,6 +89,83 @@ class QueryTraversalTest extends Specification {
 
     }
 
+    def "works for mutations()"() {
+        given:
+        def schema = TestUtil.schema("""
+            type Query {
+              a: String
+            }
+            type Mutation{
+                foo: Foo
+                bar: String
+            }
+            type Foo {
+                subFoo: String  
+            }
+            schema {mutation: Mutation, query: Query}
+        """)
+        def visitor = Mock(QueryVisitor)
+        def query = createQuery("""
+            mutation M{bar foo { subFoo} }
+            """)
+        QueryTraversal queryTraversal = createQueryTraversal(query, schema)
+        when:
+        queryTraversal."$visitFn"(visitor)
+
+        then:
+        1 * visitor.visitField({ QueryVisitorEnvironment it -> it.field.name == "foo" && it.fieldDefinition.type.name == "Foo" && it.parentType.name == "Mutation" })
+        1 * visitor.visitField({ QueryVisitorEnvironment it -> it.field.name == "bar" && it.fieldDefinition.type.name == "String" && it.parentType.name == "Mutation" })
+        1 * visitor.visitField({ QueryVisitorEnvironment it ->
+            it.field.name == "subFoo" && it.fieldDefinition.type.name == "String" &&
+                    it.parentType.name == "Foo" &&
+                    it.parentEnvironment.field.name == "foo" && it.parentEnvironment.fieldDefinition.type.name == "Foo"
+        })
+
+        where:
+        order       | visitFn
+        'postOrder' | 'visitPostOrder'
+        'preOrder'  | 'visitPreOrder'
+
+    }
+
+    def "works for subscriptions()"() {
+        given:
+        def schema = TestUtil.schema("""
+            type Query {
+              a: String
+            }
+            type Subscription{
+                foo: Foo
+                bar: String
+            }
+            type Foo {
+                subFoo: String  
+            }
+            schema {subscription: Subscription, query: Query}
+        """)
+        def visitor = Mock(QueryVisitor)
+        def query = createQuery("""
+            subscription S{bar foo { subFoo} }
+            """)
+        QueryTraversal queryTraversal = createQueryTraversal(query, schema)
+        when:
+        queryTraversal."$visitFn"(visitor)
+
+        then:
+        1 * visitor.visitField({ QueryVisitorEnvironment it -> it.field.name == "foo" && it.fieldDefinition.type.name == "Foo" && it.parentType.name == "Subscription" })
+        1 * visitor.visitField({ QueryVisitorEnvironment it -> it.field.name == "bar" && it.fieldDefinition.type.name == "String" && it.parentType.name == "Subscription" })
+        1 * visitor.visitField({ QueryVisitorEnvironment it ->
+            it.field.name == "subFoo" && it.fieldDefinition.type.name == "String" &&
+                    it.parentType.name == "Foo" &&
+                    it.parentEnvironment.field.name == "foo" && it.parentEnvironment.fieldDefinition.type.name == "Foo"
+        })
+
+        where:
+        order       | visitFn
+        'postOrder' | 'visitPostOrder'
+        'preOrder'  | 'visitPreOrder'
+
+    }
 
     @Unroll
     def "field with arguments: (#order)"() {
