Eliminate duplicated code to improve safety

Author: holly-cummins

test-framework/junit5/src/main/java/io/quarkus/test/junit/launcher/ExecutionListener.java

@@ -3,6 +3,7 @@
 import java.util.ArrayDeque;
 import java.util.Deque;
 import java.util.Optional;
+import java.util.function.Consumer;
 
 import org.junit.platform.engine.TestExecutionResult;
 import org.junit.platform.engine.TestSource;
@@ -23,42 +24,43 @@ public class ExecutionListener implements TestExecutionListener {
 
     @Override
     public void executionStarted(TestIdentifier testIdentifier) {
-        // This will be called for various levels of containers, only some of which are tests, so check carefully and do not assume
-        Optional<TestSource> oSource = testIdentifier.getSource();
-        if (oSource.isPresent()) {
-            TestSource source = oSource.get();
-            if (source instanceof ClassSource cs) {
-                ClassLoader classLoader = cs.getJavaClass().getClassLoader();
-                // Only adjust the TCCL in cases where we know the QuarkusTestExtension would be about to do it anyway
-                // We could check annotations, but that would be slow, and the assumption that only Quarkus Tests are loaded with the quarkus classloader should be a fair one
-                if (isQuarkusTest(classLoader)) {
-                    origCl.push(Thread.currentThread().getContextClassLoader());
-                    Thread.currentThread().setContextClassLoader(classLoader);
-                }
-            }
-        }
+        invokeIfTestIsQuarkusTest(testIdentifier, this::setTCCL);
     }
 
     @Override
     public void executionFinished(TestIdentifier testIdentifier, TestExecutionResult result) {
+        invokeIfTestIsQuarkusTest(testIdentifier, this::unsetTCCL);
+    }
+
+    private void invokeIfTestIsQuarkusTest(TestIdentifier testIdentifier, Consumer<ClassLoader> consumer) {
+        // This will be called for various levels of containers, only some of which are tests, so check carefully and do not assume
         Optional<TestSource> oSource = testIdentifier.getSource();
         if (oSource.isPresent()) {
             TestSource source = oSource.get();
             if (source instanceof ClassSource cs) {
                 ClassLoader classLoader = cs.getJavaClass().getClassLoader();
-                // Only pop if we meet the conditions under which we pushed
+                // Only adjust the TCCL in cases where we know the QuarkusTestExtension would be about to do it anyway
                 if (isQuarkusTest(classLoader)) {
-                    ClassLoader cl = origCl.pop();
-                    if (cl != null) {
-                        // If execution is parallel this could produce odd results, but if execution is parallel any kind of TCCL manipulation will be ill-fated
-                        Thread.currentThread().setContextClassLoader(cl);
-                    }
+                    consumer.accept(classLoader);
                 }
             }
         }
     }
 
     private static boolean isQuarkusTest(ClassLoader classLoader) {
+        // We could check annotations, but that would be slow, and the assumption that only Quarkus Tests are loaded with the quarkus classloader should be a fair one
         return classLoader instanceof QuarkusClassLoader;
     }
+
+    private void setTCCL(ClassLoader classLoader) {
+        origCl.push(Thread.currentThread().getContextClassLoader());
+        Thread.currentThread().setContextClassLoader(classLoader);
+    }
+
+    private void unsetTCCL(ClassLoader classLoader) {
+        ClassLoader cl = origCl.pop();
+        // If execution is parallel this stack logic could produce odd results, but if execution is parallel any kind of TCCL manipulation will be ill-fated
+        Thread.currentThread().setContextClassLoader(cl);
+    }
+
 }