diff --git a/src/main/java/org/spongepowered/mij/ModLauncherExtension.java b/src/main/java/org/spongepowered/mij/ModLauncherExtension.java index 64dfd14..e2d3d32 100644 --- 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 interceptTestClassConstructor(Invocation invocation, ReflectiveInvocationContext> 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 interceptTestClassConstructor(Invocation invocation, ReflectiveI public T interceptMethod(Invocation invocation, ReflectiveInvocationContext 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 index 291227e..c5f11cd 100644 --- 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;