fix result if quite after leave

andimarek · andimarek · commit 495c97e5feed · 2018-03-17T21:14:04.000+11:00
diff --git a/src/main/java/graphql/util/Traverser.java b/src/main/java/graphql/util/Traverser.java
@@ -78,6 +78,7 @@ public TraverserResult traverse(Collection<? extends T> roots, TraverserVisitor<
                 // end-of-list marker, we are done recursing children,
                 // mark the current node as fully visited
                 TraverserContext contextForLeave = (TraverserContext) traverserState.pop();
+                currentContext = contextForLeave;
                 TraversalControl traversalControl = visitor.leave(contextForLeave);
                 assertNotNull(traversalControl, "result of leave must not be null");
                 assertTrue(CONTINUE_OR_QUIT.contains(traversalControl), "result can only return CONTINUE or QUIT");
diff --git a/src/test/groovy/graphql/util/TraverserTest.groovy b/src/test/groovy/graphql/util/TraverserTest.groovy
@@ -105,13 +105,15 @@ class TraverserTest extends Specification {
     def "quit traversal in first leave"() {
         given:
         def initialData = new ArrayList()
+        def leaveResult = "Leave result"
         def leaveCount = 0
         def visitor = [
                 enter: { TraverserContext context ->
                     context.getInitialData().add(context.thisNode().number)
                     TraversalControl.CONTINUE
                 },
                 leave: { TraverserContext context ->
+                    context.setResult(leaveResult)
                     leaveCount++
                     TraversalControl.QUIT
                 },
@@ -122,9 +124,9 @@ class TraverserTest extends Specification {
         def result = Traverser.depthFirst({ n -> n.children }, initialData).traverse(root, visitor)
 
         then:
+        result.result == leaveResult
         initialData == [0, 1, 3]
         leaveCount == 1
-
     }
 
     def "abort subtree traversal depth-first"() {
