Cleanup traversal of parent contexts

mockito · Sep 30, 2019 · 5ed8a07 · 5ed8a07
1 parent acce719
commit 5ed8a07
Showing 1 changed file with 28 additions and 26 deletions.
diff --git a/subprojects/junit-jupiter/src/main/java/org/mockito/junit/jupiter/MockitoExtension.java b/subprojects/junit-jupiter/src/main/java/org/mockito/junit/jupiter/MockitoExtension.java
@@ -5,6 +5,13 @@
 package org.mockito.junit.jupiter;
 
 
+import static org.junit.jupiter.api.extension.ExtensionContext.Namespace.create;
+import static org.junit.platform.commons.support.AnnotationSupport.findAnnotation;
+
+import java.lang.reflect.Parameter;
+import java.util.LinkedHashSet;
+import java.util.Optional;
+import java.util.Set;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
@@ -22,14 +29,6 @@
 import org.mockito.junit.MockitoJUnitRunner;
 import org.mockito.quality.Strictness;
 
-import java.lang.reflect.Parameter;
-import java.util.LinkedHashSet;
-import java.util.Optional;
-import java.util.Set;
-
-import static org.junit.jupiter.api.extension.ExtensionContext.Namespace.create;
-import static org.junit.platform.commons.support.AnnotationSupport.findAnnotation;
-
 /**
  * Extension that initializes mocks and handles strict stubbings. This extension is the JUnit Jupiter equivalent
  * of our JUnit4 {@link MockitoJUnitRunner}.
@@ -194,27 +193,30 @@ private Optional<MockitoSettings> retrieveAnnotationFromTestClasses(final Extens
     }
 
     private void collectParentTestInstances(ExtensionContext context, Set<Object> testInstances) {
-        Optional<ExtensionContext> parent = context.getParent();
-
-        boolean firstParent = true;
-        while (parent.isPresent() && parent.get() != context.getRoot()) {
-            ExtensionContext parentContext = parent.get();
-            parent = parentContext.getParent();
-
-            if (firstParent) {
-                // Ignoring first parent avoids parallel execution issues
-                // We can ignore the first parent because it has the test instance that is already in 'testInstances'
-                // See how 'testInstances' is populated
-                firstParent = false;
-                continue;
-            }
+        Optional<ExtensionContext> initialParent = context.getParent();
+
+        // Ignoring first parent avoids parallel execution issues
+        // We can ignore the first parent because it has the test instance that is already in 'testInstances'
+        // See how 'testInstances' is populated
+        initialParent.ifPresent(parent -> collectParentTestInstance(parent.getParent(), context, testInstances));
+    }
+
+    @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
+    private void collectParentTestInstance(
+        Optional<ExtensionContext> parent, ExtensionContext context,
+        Set<Object> testInstances) {
+
+        parent.ifPresent(currentParent -> {
+            if (currentParent != context.getRoot()) {
+                Object testInstance = currentParent.getStore(MOCKITO).remove(TEST_INSTANCE);
 
-            Object testInstance = parentContext.getStore(MOCKITO).remove(TEST_INSTANCE);
+                if (testInstance != null) {
+                    testInstances.add(testInstance);
+                }
 
-            if (testInstance != null) {
-                testInstances.add(testInstance);
+                collectParentTestInstance(currentParent.getParent(), context, testInstances);
             }
-        }
+        });
     }
 
     /**