1556 : Update PlayJUnitExtension.java with new implementation

framework/src/play/test/PlayJUnitExtension.java

@@ -2,20 +2,19 @@
 
 import java.io.File;
 import org.apache.commons.lang3.exception.ExceptionUtils;
-import org.junit.jupiter.api.extension.BeforeAllCallback;
-import org.junit.jupiter.api.extension.BeforeEachCallback;
-import org.junit.jupiter.api.extension.ExtensionContext;
-import org.junit.jupiter.api.extension.TestExecutionExceptionHandler;
+import org.junit.jupiter.api.extension.*;
 import play.Invoker;
 import play.Invoker.DirectInvocation;
 import play.Play;
+import java.lang.reflect.Method;
 
-public class PlayJUnitExtension implements BeforeAllCallback, BeforeEachCallback, TestExecutionExceptionHandler {
+public class PlayJUnitExtension implements BeforeAllCallback, BeforeEachCallback, AfterEachCallback, InvocationInterceptor, TestExecutionExceptionHandler {
 
     public static final String invocationType = "JUnitTest";
 
     public static boolean useCustomRunner = false;
 
+    private static final ThreadLocal<Invoker.InvocationContext> invocationContext = new ThreadLocal<>();
 
     @Override
     public void beforeAll(ExtensionContext context) throws Exception {
@@ -34,24 +33,63 @@ public void beforeAll(ExtensionContext context) throws Exception {
         TestEngine.initTest(testClass);
     }
 
-
     @Override
     public void beforeEach(ExtensionContext context) throws Exception {
         if (useCustomRunner) {
             if (!Play.started) {
                 Play.forceProd = true;
                 Play.init(new File("."), getPlayId());
             }
-            Invoker.invokeInThread(new DirectInvocation() {
-                @Override
-                public void execute() throws Exception {
-                    // Nothing here; test method will run
-                }
-                @Override
-                public Invoker.InvocationContext getInvocationContext() {
-                    return new Invoker.InvocationContext(invocationType);
+            // Store invocation context in ThreadLocal for use in interceptTestMethod
+            invocationContext.set(new Invoker.InvocationContext(invocationType));
+        }
+    }
+
+    @Override
+    public void afterEach(ExtensionContext context) throws Exception {
+        if (useCustomRunner) {
+            invocationContext.remove();
+        }
+    }
+
+    @Override
+    public void interceptTestMethod(Invocation<Void> invocation, ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext) throws Throwable {
+        if (useCustomRunner) {
+            // Run the test method within the Play thread context
+            Invoker.InvocationContext ctx = new Invoker.InvocationContext(invocationType);
+            Throwable[] testException = new Throwable[1];
+
+            // Set the Play classloader for this thread
+            ClassLoader originalLoader = Thread.currentThread().getContextClassLoader();
+            if (Play.classloader != null) {
+                Thread.currentThread().setContextClassLoader(Play.classloader);
+            }
+            try {
+                Invoker.invokeInThread(new DirectInvocation() {
+                    @Override
+                    public void execute() throws Exception {
+                        try {
+                            invocation.proceed();
+                        } catch (Throwable e) {
+                            testException[0] = e;
+                        }
+                    }
+
+                    @Override
+                    public Invoker.InvocationContext getInvocationContext() {
+                        return ctx;
+                    }
+                });
+
+                if (testException[0] != null) {
+                    throw testException[0];
                 }
-            });
+            } finally {
+                // Restore original classloader
+                Thread.currentThread().setContextClassLoader(originalLoader);
+            }
+        } else {
+            invocation.proceed();
         }
     }