Pass current ExtensionContext to TestInstancePostProcessors

Prior to this commit, TIPPs were always called with the
ClassExtensionContext. However, that only contained a test instance
when Lifecycle.PER_CLASS was in use. Now, the ExtensionContext the
instance is created for, e.g. the ExtensionContext of a test method
with the default lifecycle, is used.

Issue: #910

documentation/src/test/java/example/TestInfoDemo.java

@@ -24,7 +24,8 @@
 class TestInfoDemo {
 
 	TestInfoDemo(TestInfo testInfo) {
-		assertEquals("TestInfo Demo", testInfo.getDisplayName());
+		String displayName = testInfo.getDisplayName();
+		assertTrue(displayName.equals("TEST 1") || displayName.equals("test2()"));
 	}
 
 	@BeforeEach

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/ClassExtensionContext.java

@@ -16,7 +16,7 @@
 import java.lang.reflect.Method;
 import java.util.Optional;
 
-import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.engine.execution.AbstractExtensionContext;
 import org.junit.jupiter.engine.execution.ThrowableCollector;
 import org.junit.platform.commons.meta.API;
 import org.junit.platform.engine.EngineExecutionListener;
@@ -28,9 +28,8 @@
 public final class ClassExtensionContext extends AbstractExtensionContext<ClassTestDescriptor> {
 
 	private final ThrowableCollector throwableCollector;
-	private Object testInstance;
 
-	public ClassExtensionContext(ExtensionContext parent, EngineExecutionListener engineExecutionListener,
+	public ClassExtensionContext(AbstractExtensionContext<?> parent, EngineExecutionListener engineExecutionListener,
 			ClassTestDescriptor testDescriptor, ThrowableCollector throwableCollector) {
 
 		super(parent, engineExecutionListener, testDescriptor);
@@ -47,15 +46,6 @@ public Optional<Class<?>> getTestClass() {
 		return Optional.of(getTestDescriptor().getTestClass());
 	}
 
-	void setTestInstance(Object testInstance) {
-		this.testInstance = testInstance;
-	}
-
-	@Override
-	public Optional<Object> getTestInstance() {
-		return Optional.ofNullable(this.testInstance);
-	}
-
 	@Override
 	public Optional<Method> getTestMethod() {
 		return Optional.empty();

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/ClassTestDescriptor.java

@@ -22,7 +22,6 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
-import java.util.function.Consumer;
 import java.util.function.Function;
 
 import org.junit.jupiter.api.TestInstance;
@@ -32,6 +31,7 @@
 import org.junit.jupiter.api.extension.Extension;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.TestInstancePostProcessor;
+import org.junit.jupiter.engine.execution.AbstractExtensionContext;
 import org.junit.jupiter.engine.execution.AfterEachMethodAdapter;
 import org.junit.jupiter.engine.execution.BeforeEachMethodAdapter;
 import org.junit.jupiter.engine.execution.ExecutableInvoker;
@@ -176,29 +176,29 @@ private TestInstanceProvider testInstanceProvider(JupiterEngineExecutionContext
 		if (this.lifecycle == Lifecycle.PER_CLASS) {
 			// Eagerly load test instance for BeforeAllCallbacks, if necessary,
 			// and store the instance in the ExtensionContext.
-			Object instance = instantiateAndPostProcessTestInstance(parentExecutionContext, extensionContext, registry,
-				extensionContext::setTestInstance);
-			return childRegistry -> instance;
+			Object instance = instantiateAndPostProcessTestInstance(parentExecutionContext, extensionContext, registry);
+			return (childContext, childRegistry) -> {
+				childContext.setTestInstance(instance);
+				return instance;
+			};
 		}
 
 		// else Lifecycle.PER_METHOD
-		return childRegistry -> instantiateAndPostProcessTestInstance(parentExecutionContext, extensionContext,
-			childRegistry.orElse(registry), instance -> {
-				// no extension context update required
-			});
+		return (childContext, childRegistry) -> instantiateAndPostProcessTestInstance(parentExecutionContext,
+			childContext, childRegistry.orElse(registry));
 	}
 
 	private Object instantiateAndPostProcessTestInstance(JupiterEngineExecutionContext context,
-			ExtensionContext extensionContext, ExtensionRegistry registry, Consumer<Object> testInstanceConsumer) {
+			AbstractExtensionContext<?> extensionContext, ExtensionRegistry registry) {
 
 		Object instance = instantiateTestClass(context, registry, extensionContext);
-		testInstanceConsumer.accept(instance);
+		extensionContext.setTestInstance(instance);
 		invokeTestInstancePostProcessors(registry, extensionContext);
 		return instance;
 	}
 
 	protected Object instantiateTestClass(JupiterEngineExecutionContext parentExecutionContext,
-			ExtensionRegistry registry, ExtensionContext extensionContext) {
+			ExtensionRegistry registry, AbstractExtensionContext<?> extensionContext) {
 
 		Constructor<?> constructor = ReflectionUtils.getDeclaredConstructor(this.testClass);
 		return executableInvoker.invoke(constructor, extensionContext, registry);

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/JupiterEngineDescriptor.java

@@ -13,7 +13,7 @@
 import static org.junit.jupiter.engine.extension.ExtensionRegistry.createRegistryWithDefaultExtensions;
 import static org.junit.platform.commons.meta.API.Usage.Internal;
 
-import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.engine.execution.AbstractExtensionContext;
 import org.junit.jupiter.engine.execution.JupiterEngineExecutionContext;
 import org.junit.jupiter.engine.extension.ExtensionRegistry;
 import org.junit.platform.commons.meta.API;
@@ -36,7 +36,7 @@ public JupiterEngineDescriptor(UniqueId uniqueId) {
 	public JupiterEngineExecutionContext prepare(JupiterEngineExecutionContext context) throws Exception {
 		ExtensionRegistry extensionRegistry = createRegistryWithDefaultExtensions(context.getConfigurationParameters());
 		EngineExecutionListener executionListener = context.getExecutionListener();
-		ExtensionContext extensionContext = new JupiterEngineExtensionContext(executionListener, this);
+		AbstractExtensionContext<?> extensionContext = new JupiterEngineExtensionContext(executionListener, this);
 
 		// @formatter:off
 		return context.extend()

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/JupiterEngineExtensionContext.java

@@ -16,6 +16,7 @@
 import java.lang.reflect.Method;
 import java.util.Optional;
 
+import org.junit.jupiter.engine.execution.AbstractExtensionContext;
 import org.junit.platform.commons.meta.API;
 import org.junit.platform.engine.EngineExecutionListener;
 
@@ -40,11 +41,6 @@ public Optional<Class<?>> getTestClass() {
 		return Optional.empty();
 	}
 
-	@Override
-	public Optional<Object> getTestInstance() {
-		return Optional.empty();
-	}
-
 	@Override
 	public Optional<Method> getTestMethod() {
 		return Optional.empty();

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/MethodExtensionContext.java

@@ -16,7 +16,7 @@
 import java.lang.reflect.Method;
 import java.util.Optional;
 
-import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.engine.execution.AbstractExtensionContext;
 import org.junit.jupiter.engine.execution.ThrowableCollector;
 import org.junit.platform.commons.meta.API;
 import org.junit.platform.engine.EngineExecutionListener;
@@ -27,16 +27,13 @@
 @API(Internal)
 public final class MethodExtensionContext extends AbstractExtensionContext<MethodTestDescriptor> {
 
-	private final Object testInstance;
-
 	private final ThrowableCollector throwableCollector;
 
-	public MethodExtensionContext(ExtensionContext parent, EngineExecutionListener engineExecutionListener,
-			MethodTestDescriptor testDescriptor, Object testInstance, ThrowableCollector throwableCollector) {
+	public MethodExtensionContext(AbstractExtensionContext<?> parent, EngineExecutionListener engineExecutionListener,
+			MethodTestDescriptor testDescriptor, ThrowableCollector throwableCollector) {
 
 		super(parent, engineExecutionListener, testDescriptor);
 
-		this.testInstance = testInstance;
 		this.throwableCollector = throwableCollector;
 	}
 
@@ -55,11 +52,6 @@ public Optional<Method> getTestMethod() {
 		return Optional.of(getTestDescriptor().getTestMethod());
 	}
 
-	@Override
-	public Optional<Object> getTestInstance() {
-		return Optional.of(this.testInstance);
-	}
-
 	@Override
 	public Optional<Throwable> getExecutionException() {
 		return Optional.ofNullable(this.throwableCollector.getThrowable());

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/MethodTestDescriptor.java

@@ -25,6 +25,7 @@
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.TestExecutionExceptionHandler;
 import org.junit.jupiter.api.function.Executable;
+import org.junit.jupiter.engine.execution.AbstractExtensionContext;
 import org.junit.jupiter.engine.execution.AfterEachMethodAdapter;
 import org.junit.jupiter.engine.execution.BeforeEachMethodAdapter;
 import org.junit.jupiter.engine.execution.ExecutableInvoker;
@@ -78,10 +79,10 @@ public Type getType() {
 	@Override
 	public JupiterEngineExecutionContext prepare(JupiterEngineExecutionContext context) throws Exception {
 		ExtensionRegistry registry = populateNewExtensionRegistry(context);
-		Object testInstance = context.getTestInstanceProvider().getTestInstance(Optional.of(registry));
 		ThrowableCollector throwableCollector = new ThrowableCollector();
-		ExtensionContext extensionContext = new MethodExtensionContext(context.getExtensionContext(),
-			context.getExecutionListener(), this, testInstance, throwableCollector);
+		AbstractExtensionContext<?> extensionContext = new MethodExtensionContext(context.getExtensionContext(),
+			context.getExecutionListener(), this, throwableCollector);
+		context.getTestInstanceProvider().getTestInstance(extensionContext, Optional.of(registry));
 
 		// @formatter:off
 		return context.extend()

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/NestedClassTestDescriptor.java

@@ -16,7 +16,7 @@
 import java.util.Optional;
 import java.util.Set;
 
-import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.engine.execution.AbstractExtensionContext;
 import org.junit.jupiter.engine.execution.ExecutableInvoker;
 import org.junit.jupiter.engine.execution.JupiterEngineExecutionContext;
 import org.junit.jupiter.engine.extension.ExtensionRegistry;
@@ -58,11 +58,11 @@ public final Set<TestTag> getTags() {
 
 	@Override
 	protected Object instantiateTestClass(JupiterEngineExecutionContext parentExecutionContext,
-			ExtensionRegistry registry, ExtensionContext extensionContext) {
+			ExtensionRegistry registry, AbstractExtensionContext<?> extensionContext) {
 
 		// Extensions registered for nested classes and below are not to be used for instantiating outer classes
 		Optional<ExtensionRegistry> childExtensionRegistryForOuterInstance = Optional.empty();
-		Object outerInstance = parentExecutionContext.getTestInstanceProvider().getTestInstance(
+		Object outerInstance = parentExecutionContext.getTestInstanceProvider().getTestInstance(extensionContext,
 			childExtensionRegistryForOuterInstance);
 		Constructor<?> constructor = ReflectionUtils.getDeclaredConstructor(getTestClass());
 		return executableInvoker.invoke(constructor, outerInstance, extensionContext, registry);

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/TestTemplateExtensionContext.java

@@ -16,7 +16,7 @@
 import java.lang.reflect.Method;
 import java.util.Optional;
 
-import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.engine.execution.AbstractExtensionContext;
 import org.junit.platform.commons.meta.API;
 import org.junit.platform.engine.EngineExecutionListener;
 
@@ -26,13 +26,10 @@
 @API(Internal)
 final class TestTemplateExtensionContext extends AbstractExtensionContext<TestTemplateTestDescriptor> {
 
-	private final Object testInstance;
-
-	TestTemplateExtensionContext(ExtensionContext parent, EngineExecutionListener engineExecutionListener,
-			TestTemplateTestDescriptor testDescriptor, Object testInstance) {
+	TestTemplateExtensionContext(AbstractExtensionContext<?> parent, EngineExecutionListener engineExecutionListener,
+			TestTemplateTestDescriptor testDescriptor) {
 
 		super(parent, engineExecutionListener, testDescriptor);
-		this.testInstance = testInstance;
 	}
 
 	@Override
@@ -45,11 +42,6 @@ public Optional<Class<?>> getTestClass() {
 		return Optional.of(getTestDescriptor().getTestClass());
 	}
 
-	@Override
-	public Optional<Object> getTestInstance() {
-		return Optional.ofNullable(this.testInstance);
-	}
-
 	@Override
 	public Optional<Method> getTestMethod() {
 		return Optional.of(getTestDescriptor().getTestMethod());

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/TestTemplateTestDescriptor.java

@@ -20,9 +20,9 @@
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.TestTemplateInvocationContext;
 import org.junit.jupiter.api.extension.TestTemplateInvocationContextProvider;
+import org.junit.jupiter.engine.execution.AbstractExtensionContext;
 import org.junit.jupiter.engine.execution.JupiterEngineExecutionContext;
 import org.junit.jupiter.engine.extension.ExtensionRegistry;
-import org.junit.platform.commons.JUnitException;
 import org.junit.platform.commons.meta.API;
 import org.junit.platform.commons.util.Preconditions;
 import org.junit.platform.engine.TestDescriptor;
@@ -59,12 +59,11 @@ public JupiterEngineExecutionContext prepare(JupiterEngineExecutionContext conte
 		ExtensionRegistry registry = populateNewExtensionRegistryFromExtendWith(getTestMethod(),
 			context.getExtensionRegistry());
 
-		// The test instance should be properly maintained by the enclosing class's ExtensionContext.
-		Object testInstance = context.getExtensionContext().getTestInstance().orElseThrow(() -> new JUnitException(
-			"Illegal state: test instance not present for method: " + getTestMethod().toGenericString()));
-
-		ExtensionContext extensionContext = new TestTemplateExtensionContext(context.getExtensionContext(),
-			context.getExecutionListener(), this, testInstance);
+		AbstractExtensionContext<?> extensionContext = new TestTemplateExtensionContext(context.getExtensionContext(),
+			context.getExecutionListener(), this);
+		// We don't require a test instance here but provide it to extensions should the enclosing class's
+		// ExtensionContext already contain one.
+		extensionContext.setTestInstance(context.getExtensionContext().getTestInstance().orElse(null));
 
 		// @formatter:off
 		return context.extend()

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/AbstractExtensionContext.java → junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/execution/AbstractExtensionContext.java

@@ -8,18 +8,18 @@
  * http://www.eclipse.org/legal/epl-v10.html
  */
 
-package org.junit.jupiter.engine.descriptor;
+package org.junit.jupiter.engine.execution;
 
 import static java.util.stream.Collectors.toCollection;
+import static org.junit.platform.commons.meta.API.Usage.Internal;
 
 import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
 import org.junit.jupiter.api.extension.ExtensionContext;
-import org.junit.jupiter.engine.execution.ExtensionValuesStore;
-import org.junit.jupiter.engine.execution.NamespaceAwareStore;
+import org.junit.platform.commons.meta.API;
 import org.junit.platform.commons.util.Preconditions;
 import org.junit.platform.engine.EngineExecutionListener;
 import org.junit.platform.engine.TestDescriptor;
@@ -29,25 +29,28 @@
 /**
  * @since 5.0
  */
-abstract class AbstractExtensionContext<T extends TestDescriptor> implements ExtensionContext {
+@API(Internal)
+public abstract class AbstractExtensionContext<T extends TestDescriptor> implements ExtensionContext {
 
 	private final ExtensionContext parent;
 	private final EngineExecutionListener engineExecutionListener;
 	private final T testDescriptor;
 	private final ExtensionValuesStore valuesStore;
 
-	AbstractExtensionContext(ExtensionContext parent, EngineExecutionListener engineExecutionListener,
-			T testDescriptor) {
+	private Object testInstance;
+
+	protected AbstractExtensionContext(AbstractExtensionContext<?> parent,
+			EngineExecutionListener engineExecutionListener, T testDescriptor) {
 		this.parent = parent;
 		this.engineExecutionListener = engineExecutionListener;
 		this.testDescriptor = testDescriptor;
 		this.valuesStore = createStore(parent);
 	}
 
-	private ExtensionValuesStore createStore(ExtensionContext parent) {
+	private ExtensionValuesStore createStore(AbstractExtensionContext<?> parent) {
 		ExtensionValuesStore parentStore = null;
 		if (parent != null) {
-			parentStore = ((AbstractExtensionContext<?>) parent).valuesStore;
+			parentStore = parent.valuesStore;
 		}
 		return new ExtensionValuesStore(parentStore);
 	}
@@ -62,6 +65,15 @@ public String getDisplayName() {
 		return getTestDescriptor().getDisplayName();
 	}
 
+	public void setTestInstance(Object testInstance) {
+		this.testInstance = testInstance;
+	}
+
+	@Override
+	public Optional<Object> getTestInstance() {
+		return Optional.ofNullable(this.testInstance);
+	}
+
 	@Override
 	public void publishReportEntry(Map<String, String> values) {
 		engineExecutionListener.reportingEntryPublished(this.testDescriptor, ReportEntry.from(values));

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/execution/JupiterEngineExecutionContext.java

@@ -12,7 +12,6 @@
 
 import static org.junit.platform.commons.meta.API.Usage.Internal;
 
-import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.engine.extension.ExtensionRegistry;
 import org.junit.platform.commons.JUnitException;
 import org.junit.platform.commons.meta.API;
@@ -56,7 +55,7 @@ public ExtensionRegistry getExtensionRegistry() {
 		return this.state.extensionRegistry;
 	}
 
-	public ExtensionContext getExtensionContext() {
+	public AbstractExtensionContext<?> getExtensionContext() {
 		return this.state.extensionContext;
 	}
 
@@ -82,7 +81,7 @@ private static final class State implements Cloneable {
 		final ConfigurationParameters configurationParameters;
 		TestInstanceProvider testInstanceProvider;
 		ExtensionRegistry extensionRegistry;
-		ExtensionContext extensionContext;
+		AbstractExtensionContext<?> extensionContext;
 		ThrowableCollector throwableCollector;
 
 		State(EngineExecutionListener executionListener, ConfigurationParameters configurationParameters) {
@@ -121,7 +120,7 @@ public Builder withExtensionRegistry(ExtensionRegistry extensionRegistry) {
 			return this;
 		}
 
-		public Builder withExtensionContext(ExtensionContext extensionContext) {
+		public Builder withExtensionContext(AbstractExtensionContext<?> extensionContext) {
 			newState().extensionContext = extensionContext;
 			return this;
 		}