fixes: fix tests

Author: yarinvak

src/main/java/graphql/annotations/AnnotationsSchemaCreator.java

@@ -16,6 +16,7 @@
 
 import graphql.annotations.directives.AnnotationsDirectiveWiring;
 import graphql.annotations.directives.DirectiveSchemaVisitor;
+import graphql.annotations.directives.TreeTransformerUtilWrapper;
 import graphql.annotations.processor.DirectiveAndWiring;
 import graphql.annotations.processor.GraphQLAnnotations;
 import graphql.annotations.processor.typeFunctions.TypeFunction;
@@ -267,9 +268,10 @@ public GraphQLSchema build() {
             GraphQLSchema schema = this.graphqlSchemaBuilder.build();
             // wire with directives
             HashMap<String, AnnotationsDirectiveWiring> directiveWiringHashMap = transformDirectiveRegistry(this.graphQLAnnotations.getContainer().getDirectiveRegistry());
-            DirectiveSchemaVisitor directiveSchemaVisitor = new DirectiveSchemaVisitor(directiveWiringHashMap, graphQLAnnotations.getContainer().getCodeRegistryBuilder());
-            GraphQLSchema trasnformedSchema = this.schemaTransformer.transform(schema, directiveSchemaVisitor);
-            return newSchema(trasnformedSchema).codeRegistry(graphQLAnnotations.getContainer().getCodeRegistryBuilder().build()).build();
+            DirectiveSchemaVisitor directiveSchemaVisitor = new DirectiveSchemaVisitor(directiveWiringHashMap,
+                    graphQLAnnotations.getContainer().getCodeRegistryBuilder(), new TreeTransformerUtilWrapper());
+            GraphQLSchema transformedSchema = this.schemaTransformer.transform(schema, directiveSchemaVisitor);
+            return newSchema(transformedSchema).codeRegistry(graphQLAnnotations.getContainer().getCodeRegistryBuilder().build()).build();
         }
 
         private HashMap<String, AnnotationsDirectiveWiring> transformDirectiveRegistry(Map<String, DirectiveAndWiring> directiveRegistry) {

src/main/java/graphql/annotations/directives/DirectiveSchemaVisitor.java

@@ -11,11 +11,10 @@
 import java.util.List;
 import java.util.Map;
 
-import static graphql.util.TreeTransformerUtil.changeNode;
-
 public class DirectiveSchemaVisitor implements GraphQLTypeVisitor {
     private HashMap<String, AnnotationsDirectiveWiring> directiveWiringMap;
     private GraphQLCodeRegistry.Builder codeRegistryBuilder;
+    private TreeTransformerUtilWrapper transformerUtilWrapper;
 
     @FunctionalInterface
     interface WiringFunction {
@@ -27,18 +26,19 @@ GraphQLDirectiveContainer apply(GraphQLDirective a, GraphQLDirectiveContainer b,
     private Map<Class, WiringFunction> functionMap;
 
 
-    public DirectiveSchemaVisitor(HashMap<String, AnnotationsDirectiveWiring> directiveWiringMap, GraphQLCodeRegistry.Builder codeRegistryBuilder) {
+    public DirectiveSchemaVisitor(HashMap<String, AnnotationsDirectiveWiring> directiveWiringMap, GraphQLCodeRegistry.Builder codeRegistryBuilder,
+                                  TreeTransformerUtilWrapper treeTransformerUtilWrapper) {
         this.directiveWiringMap = directiveWiringMap;
         this.functionMap = createFunctionsMap();
         this.codeRegistryBuilder = codeRegistryBuilder;
+        this.transformerUtilWrapper = treeTransformerUtilWrapper;
     }
 
     @Override
     public TraversalControl visitGraphQLArgument(GraphQLArgument node, TraverserContext<GraphQLSchemaElement> context) {
         return this.visitGraphQLType(GraphQLArgument.class, node, context);
     }
 
-
     @Override
     public TraversalControl visitGraphQLInterfaceType(GraphQLInterfaceType node, TraverserContext<GraphQLSchemaElement> context) {
         return this.visitGraphQLType(GraphQLInterfaceType.class, node, context);
@@ -124,7 +124,7 @@ private TraversalControl visitGraphQLType(Class<? extends GraphQLDirectiveContai
                 }
             }
         }
-        return changeNode(context, newNode);
+        return transformerUtilWrapper.changeNode(context, newNode);
     }
 
     private void putInMap(Map<Class, WiringFunction> map, Class clazz, String functionName,

src/main/java/graphql/annotations/directives/TreeTransformerUtilWrapper.java

@@ -0,0 +1,11 @@
+package graphql.annotations.directives;
+
+import graphql.util.TraversalControl;
+import graphql.util.TraverserContext;
+import graphql.util.TreeTransformerUtil;
+
+public class TreeTransformerUtilWrapper {
+    public <T> TraversalControl changeNode(TraverserContext<T> context, T changedNode) {
+        return TreeTransformerUtil.changeNode(context, changedNode);
+    }
+}