test for batch limiting

Author: jmccaull

src/test/groovy/graphql/servlet/AbstractGraphQLHttpServletSpec.groovy

@@ -285,6 +285,19 @@ class AbstractGraphQLHttpServletSpec extends Specification {
         getBatchedResponseContent()[1].data.echo == "test"
     }
 
+    def "Execution Result Handler allows limiting number of queries"() {
+        setup:
+        request.addParameter('query', '[{ "query": "query { echo(arg:\\"test\\") }" }, { "query": "query { echo(arg:\\"test\\") }" }, { "query": "query { echo(arg:\\"test\\") }" }]')
+
+        when:
+        servlet.doGet(request, response)
+
+        then:
+        response.getStatus() == STATUS_OK
+        response.getContentType() == CONTENT_TYPE_JSON_UTF8
+        getBatchedResponseContent().size() == 2
+    }
+
     def "mutation over HTTP GET returns errors"() {
         setup:
         request.addParameter('query', 'mutation { echo(arg:"test") }')

src/test/groovy/graphql/servlet/TestBatchInputHandlerFactory.java

@@ -0,0 +1,56 @@
+package graphql.servlet;
+
+import graphql.ExecutionInput;
+import graphql.ExecutionResult;
+
+import java.io.IOException;
+import java.io.Writer;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+public class TestBatchInputHandlerFactory implements GraphQLExecutionResultHandlerFactory {
+    @Override
+    public ExecutionResultHandler getBatchHandler(Writer respWriter, GraphQLObjectMapper graphQLObjectMapper) {
+        return new LimitedBatchSizeHandler(respWriter, graphQLObjectMapper);
+    }
+
+    private class LimitedBatchSizeHandler implements ExecutionResultHandler {
+
+        Writer writer;
+        GraphQLObjectMapper mapper;
+        List<ExecutionResult> results = new ArrayList<>();
+
+        private LimitedBatchSizeHandler(Writer respWriter, GraphQLObjectMapper graphQLObjectMapper) {
+            this.mapper = graphQLObjectMapper;
+            this.writer = respWriter;
+        }
+
+        @Override
+        public void finalizeResults() {
+            try {
+                writer.write("[");
+                Iterator<ExecutionResult> iter = results.iterator();
+                while (iter.hasNext()) {
+                    writer.write(mapper.serializeResultAsJson(iter.next()));
+                    if (iter.hasNext()) {
+                        writer.write(",");
+                    }
+                }
+                writer.write("]");
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        @Override
+        public void accept(ExecutionResult executionResult) {
+            results.add(executionResult);
+        }
+
+        @Override
+        public boolean shouldContinue(Iterator<ExecutionInput> executionInputIterator) {
+            return results.size() < 2 && executionInputIterator.hasNext();
+        }
+    }
+}

src/test/groovy/graphql/servlet/TestUtils.groovy

@@ -24,6 +24,7 @@ class TestUtils {
                 .with(createGraphQlSchema(queryDataFetcher, mutationDataFetcher, subscriptionDataFetcher))
                 .with(createInstrumentedQueryInvoker())
                 .with(asyncServletModeEnabled)
+                .with(createBatchInputHandlerFactory())
                 .build())
         servlet.init(null)
         return servlet
@@ -34,6 +35,10 @@ class TestUtils {
         GraphQLQueryInvoker.newBuilder().with([instrumentation]).build()
     }
 
+    static def createBatchInputHandlerFactory() {
+        new TestBatchInputHandlerFactory()
+    }
+
     static def createGraphQlSchema(DataFetcher queryDataFetcher = { env -> env.arguments.arg },
                                    DataFetcher mutationDataFetcher = { env -> env.arguments.arg },
                                    DataFetcher subscriptionDataFetcher = { env ->