relay tests fixed

Author: yarinvak

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

@@ -82,6 +82,15 @@ public Builder typeFunction(TypeFunction typeFunction) {
             return this;
         }
 
+        public Builder setAlwaysPrettify(Boolean shouldAlwaysPrettify) {
+            this.graphQLAnnotations.getObjectHandler().getTypeRetriever().getGraphQLFieldRetriever().setAlwaysPrettify(shouldAlwaysPrettify);
+            return this;
+        }
+
+        public GraphQLAnnotations getGraphQLAnnotations() {
+            return this.graphQLAnnotations;
+        }
+
         public GraphQLSchema build() {
             assert this.queryObject != null;
 

src/main/java/graphql/annotations/processor/retrievers/GraphQLFieldRetriever.java

@@ -31,6 +31,7 @@
 import graphql.annotations.processor.retrievers.fieldBuilders.method.MethodNameBuilder;
 import graphql.annotations.processor.retrievers.fieldBuilders.method.MethodTypeBuilder;
 import graphql.annotations.processor.typeFunctions.TypeFunction;
+import graphql.annotations.processor.util.CodeRegistryUtil;
 import graphql.annotations.processor.util.ConnectionUtil;
 import graphql.annotations.processor.util.DataFetcherConstructor;
 import graphql.relay.Relay;
@@ -151,6 +152,14 @@ private GraphQLFieldDefinition handleRelayArguments(Method method, ProcessingEle
         GraphQLFieldDefinition relayFieldDefinition = null;
         if (method.isAnnotationPresent(GraphQLRelayMutation.class)) {
             relayFieldDefinition = buildRelayMutation(method, container, builder, outputType, args);
+
+            // Getting the data fetcher from the old field type and putting it as the new type
+            String newParentType = relayFieldDefinition.getType().getName();
+            relayFieldDefinition.getType().getChildren().forEach(field -> {
+                DataFetcher dataFetcher = CodeRegistryUtil.getDataFetcher(container.getCodeRegistryBuilder(), outputType.getName(), (GraphQLFieldDefinition) field);
+                container.getCodeRegistryBuilder().dataFetcher(coordinates(newParentType, field.getName()), dataFetcher);
+            });
+
         } else {
             builder.arguments(args);
         }

src/test/java/graphql/annotations/RelayTest.java

@@ -30,6 +30,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import static graphql.annotations.AnnotationsSchemaCreator.newAnnotationsSchema;
 import static org.testng.Assert.*;
 
 @SuppressWarnings("unchecked")
@@ -69,7 +70,6 @@ public Result(int i) {
             this.i = i;
         }
 
-        @GraphQLField
         public int getI() {
             return i;
         }
@@ -129,7 +129,7 @@ public void noArgMutation() {
         assertNotNull(returnType.getFieldDefinition("getI"));
         assertNotNull(returnType.getFieldDefinition("clientMutationId"));
 
-        GraphQLSchema schema = GraphQLSchema.newSchema().query(object).mutation(object).build();
+        GraphQLSchema schema = newAnnotationsSchema().query(TestObject.class).mutation(TestObject.class).build();
 
         GraphQL graphQL = GraphQL.newGraphQL(schema).queryExecutionStrategy(new EnhancedExecutionStrategy()).build();
 
@@ -151,7 +151,7 @@ public void interfaceReturningMutation() {
 
         assertNotNull(doSomething);
 
-        GraphQLSchema schema = GraphQLSchema.newSchema().query(object).mutation(object).build();
+        GraphQLSchema schema = newAnnotationsSchema().query(TestObject.class).mutation(TestObject.class).build();
 
         GraphQL graphQL = GraphQL.newGraphQL(schema).queryExecutionStrategy(new EnhancedExecutionStrategy()).build();
 
@@ -190,7 +190,7 @@ public void argMutation() {
         assertNotNull(returnType.getFieldDefinition("getI"));
         assertNotNull(returnType.getFieldDefinition("clientMutationId"));
 
-        GraphQLSchema schema = GraphQLSchema.newSchema().query(object).mutation(object).codeRegistry(this.graphQLAnnotations.getContainer().getCodeRegistryBuilder().build()).build();
+        GraphQLSchema schema = newAnnotationsSchema().query(TestObject.class).mutation(TestObject.class).build();
 
         GraphQL graphQL = GraphQL.newGraphQL(schema).queryExecutionStrategy(new EnhancedExecutionStrategy()).build();
 
@@ -227,7 +227,7 @@ public void argVariableMutation() {
         assertNotNull(returnType.getFieldDefinition("getI"));
         assertNotNull(returnType.getFieldDefinition("clientMutationId"));
 
-        GraphQLSchema schema = GraphQLSchema.newSchema().query(object).mutation(object).build();
+        GraphQLSchema schema = newAnnotationsSchema().query(TestObject.class).mutation(TestObject.class).build();
 
         GraphQL graphQL = GraphQL.newGraphQL(schema).queryExecutionStrategy(new EnhancedExecutionStrategy()).build();