Set ContextClassLoader when invoking tests

SpongePowered · Jan 13, 2023 · d7e0d4b · d7e0d4b
1 parent aaa39b4
commit d7e0d4b
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 10 deletions.
diff --git a/src/main/java/org/spongepowered/mij/ModLauncherExtension.java b/src/main/java/org/spongepowered/mij/ModLauncherExtension.java
@@ -44,8 +44,18 @@ public abstract class ModLauncherExtension implements MethodInvocationIntercepto
     public <T> T interceptTestClassConstructor(Invocation<T> invocation, ReflectiveInvocationContext<Constructor<T>> invocationContext,
             ExtensionContext extensionContext) throws Throwable {
         final T originalInstance = invocation.proceed();
-        final Object transformedInstance = ReflectionUtils.newInstance(getTransformedConstructor(invocationContext.getExecutable()),
-                invocationContext.getArguments().toArray());
+
+        final Thread thread = Thread.currentThread();
+        final ClassLoader originalClassLoader = thread.getContextClassLoader();
+
+        final Object transformedInstance;
+        try {
+            thread.setContextClassLoader(getTransformingClassLoader());
+            transformedInstance = ReflectionUtils.newInstance(getTransformedConstructor(invocationContext.getExecutable()),
+                    invocationContext.getArguments().toArray());
+        } finally {
+            thread.setContextClassLoader(originalClassLoader);
+        }
 
         this.originalToTransformedInstances.put(originalInstance, transformedInstance);
         return originalInstance;
@@ -56,8 +66,17 @@ public <T> T interceptTestClassConstructor(Invocation<T> invocation, ReflectiveI
     public <T> T interceptMethod(Invocation<T> invocation, ReflectiveInvocationContext<Method> invocationContext,
             ExtensionContext extensionContext) throws Throwable {
         invocation.skip();
-        return (T) ReflectionUtils.invokeMethod(getTransformedMethod(invocationContext.getExecutable()),
-                getTransformedInstance(invocationContext.getTarget().orElse(null)), invocationContext.getArguments().toArray());
+
+        final Thread thread = Thread.currentThread();
+        final ClassLoader originalClassLoader = thread.getContextClassLoader();
+
+        try {
+            thread.setContextClassLoader(getTransformingClassLoader());
+            return (T) ReflectionUtils.invokeMethod(getTransformedMethod(invocationContext.getExecutable()),
+                    getTransformedInstance(invocationContext.getTarget().orElse(null)), invocationContext.getArguments().toArray());
+        } finally {
+            thread.setContextClassLoader(originalClassLoader);
+        }
     }
 
     protected Object getTransformedInstance(Object originalInstance) {

diff --git a/src/main/java/org/spongepowered/mij/SharedModLauncher.java b/src/main/java/org/spongepowered/mij/SharedModLauncher.java
@@ -46,13 +46,15 @@ public static ClassLoader getTransformingClassLoader(String[] launcherArgs) {
             final Thread thread = Thread.currentThread();
             final ClassLoader originalClassLoader = thread.getContextClassLoader();
 
-            Launcher.main(launcherArgs);
+            try {
+                Launcher.main(launcherArgs);
 
-            // Capture the context class loader set by ModLauncher
-            SharedModLauncher.transformingClassLoader = thread.getContextClassLoader();
-
-            // Restore the original context class loader to avoid potential issues with classic tests
-            thread.setContextClassLoader(originalClassLoader);
+                // Capture the context class loader set by ModLauncher
+                SharedModLauncher.transformingClassLoader = thread.getContextClassLoader();
+            } finally {
+                // Restore the original context class loader to avoid potential issues with classic tests
+                thread.setContextClassLoader(originalClassLoader);
+            }
         }
 
         return SharedModLauncher.transformingClassLoader;