Fix ExtensionContext on instance creation

Issue: junit-team#3445

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/ClassBasedTestDescriptor.java

@@ -201,8 +201,7 @@ public JupiterEngineExecutionContext before(JupiterEngineExecutionContext contex
 			// and store the instance in the ExtensionContext.
 			ClassExtensionContext extensionContext = (ClassExtensionContext) context.getExtensionContext();
 			throwableCollector.execute(() -> {
-				TestInstances testInstances = context.getTestInstancesProvider().getTestInstances(
-					context.getExtensionRegistry(), throwableCollector);
+				TestInstances testInstances = context.getTestInstancesProvider().getTestInstances(context);
 				extensionContext.setTestInstances(testInstances);
 			});
 		}
@@ -273,32 +272,30 @@ private TestInstanceFactory resolveTestInstanceFactory(ExtensionRegistry registr
 	}
 
 	private TestInstancesProvider testInstancesProvider(JupiterEngineExecutionContext parentExecutionContext,
-			ClassExtensionContext extensionContext) {
+			ClassExtensionContext ourExtensionContext) {
 
-		return (registry, registrar, throwableCollector) -> extensionContext.getTestInstances().orElseGet(
-			() -> instantiateAndPostProcessTestInstance(parentExecutionContext, extensionContext, registry, registrar,
-				throwableCollector));
+		// For Lifecycle.PER_CLASS, ourExtensionContext.getTestInstances() is used to store the instance.
+		// Otherwise, extensionContext.getTestInstances() is always empty and we always create a new instance.
+		return (registry, context) -> ourExtensionContext.getTestInstances().orElseGet(
+			() -> instantiateAndPostProcessTestInstance(parentExecutionContext, registry, context));
 	}
 
 	private TestInstances instantiateAndPostProcessTestInstance(JupiterEngineExecutionContext parentExecutionContext,
-			ExtensionContext extensionContext, ExtensionRegistry registry, ExtensionRegistrar registrar,
-			ThrowableCollector throwableCollector) {
+			ExtensionRegistry registry, JupiterEngineExecutionContext context) {
 
-		TestInstances instances = instantiateTestClass(parentExecutionContext, registry, registrar, extensionContext,
-			throwableCollector);
-		throwableCollector.execute(() -> {
-			invokeTestInstancePostProcessors(instances.getInnermostInstance(), registry, extensionContext);
+		TestInstances instances = instantiateTestClass(parentExecutionContext, registry, context);
+		context.getThrowableCollector().execute(() -> {
+			invokeTestInstancePostProcessors(instances.getInnermostInstance(), registry, context.getExtensionContext());
 			// In addition, we initialize extension registered programmatically from instance fields here
 			// since the best time to do that is immediately following test class instantiation
 			// and post-processing.
-			registrar.initializeExtensions(this.testClass, instances.getInnermostInstance());
+			context.getExtensionRegistry().initializeExtensions(this.testClass, instances.getInnermostInstance());
 		});
 		return instances;
 	}
 
 	protected abstract TestInstances instantiateTestClass(JupiterEngineExecutionContext parentExecutionContext,
-			ExtensionRegistry registry, ExtensionRegistrar registrar, ExtensionContext extensionContext,
-			ThrowableCollector throwableCollector);
+			ExtensionRegistry registry, JupiterEngineExecutionContext context);
 
 	protected TestInstances instantiateTestClass(Optional<TestInstances> outerInstances, ExtensionRegistry registry,
 			ExtensionContext extensionContext) {

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

@@ -20,16 +20,13 @@
 import java.util.Set;
 
 import org.apiguardian.api.API;
-import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.TestInstances;
 import org.junit.jupiter.engine.config.JupiterConfiguration;
 import org.junit.jupiter.engine.execution.JupiterEngineExecutionContext;
-import org.junit.jupiter.engine.extension.ExtensionRegistrar;
 import org.junit.jupiter.engine.extension.ExtensionRegistry;
 import org.junit.platform.engine.TestDescriptor;
 import org.junit.platform.engine.TestTag;
 import org.junit.platform.engine.UniqueId;
-import org.junit.platform.engine.support.hierarchical.ThrowableCollector;
 
 /**
  * {@link TestDescriptor} for tests based on Java classes.
@@ -74,9 +71,8 @@ public ExecutionMode getExecutionMode() {
 
 	@Override
 	protected TestInstances instantiateTestClass(JupiterEngineExecutionContext parentExecutionContext,
-			ExtensionRegistry registry, ExtensionRegistrar registrar, ExtensionContext extensionContext,
-			ThrowableCollector throwableCollector) {
-		return instantiateTestClass(Optional.empty(), registry, extensionContext);
+			ExtensionRegistry registry, JupiterEngineExecutionContext context) {
+		return instantiateTestClass(Optional.empty(), registry, context.getExtensionContext());
 	}
 
 }

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

@@ -21,16 +21,13 @@
 import java.util.Set;
 
 import org.apiguardian.api.API;
-import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.TestInstances;
 import org.junit.jupiter.engine.config.JupiterConfiguration;
 import org.junit.jupiter.engine.execution.JupiterEngineExecutionContext;
-import org.junit.jupiter.engine.extension.ExtensionRegistrar;
 import org.junit.jupiter.engine.extension.ExtensionRegistry;
 import org.junit.platform.engine.TestDescriptor;
 import org.junit.platform.engine.TestTag;
 import org.junit.platform.engine.UniqueId;
-import org.junit.platform.engine.support.hierarchical.ThrowableCollector;
 
 /**
  * {@link TestDescriptor} for tests based on nested (but not static) Java classes.
@@ -77,14 +74,13 @@ public List<Class<?>> getEnclosingTestClasses() {
 
 	@Override
 	protected TestInstances instantiateTestClass(JupiterEngineExecutionContext parentExecutionContext,
-			ExtensionRegistry registry, ExtensionRegistrar registrar, ExtensionContext extensionContext,
-			ThrowableCollector throwableCollector) {
+			ExtensionRegistry registry, JupiterEngineExecutionContext context) {
 
 		// Extensions registered for nested classes and below are not to be used for instantiating and initializing outer classes
 		ExtensionRegistry extensionRegistryForOuterInstanceCreation = parentExecutionContext.getExtensionRegistry();
 		TestInstances outerInstances = parentExecutionContext.getTestInstancesProvider().getTestInstances(
-			extensionRegistryForOuterInstanceCreation, registrar, throwableCollector);
-		return instantiateTestClass(Optional.of(outerInstances), registry, extensionContext);
+			extensionRegistryForOuterInstanceCreation, context);
+		return instantiateTestClass(Optional.of(outerInstances), registry, context.getExtensionContext());
 	}
 
 }

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

@@ -101,19 +101,18 @@ public JupiterEngineExecutionContext prepare(JupiterEngineExecutionContext conte
 		MethodExtensionContext extensionContext = new MethodExtensionContext(context.getExtensionContext(),
 			context.getExecutionListener(), this, context.getConfiguration(), throwableCollector,
 			it -> new DefaultExecutableInvoker(it, registry));
-		throwableCollector.execute(() -> {
-			TestInstances testInstances = context.getTestInstancesProvider().getTestInstances(registry,
-				throwableCollector);
-			extensionContext.setTestInstances(testInstances);
-		});
-
 		// @formatter:off
-		return context.extend()
+		JupiterEngineExecutionContext newContext = context.extend()
 				.withExtensionRegistry(registry)
 				.withExtensionContext(extensionContext)
 				.withThrowableCollector(throwableCollector)
 				.build();
 		// @formatter:on
+		throwableCollector.execute(() -> {
+			TestInstances testInstances = newContext.getTestInstancesProvider().getTestInstances(newContext);
+			extensionContext.setTestInstances(testInstances);
+		});
+		return newContext;
 	}
 
 	protected MutableExtensionRegistry populateNewExtensionRegistry(JupiterEngineExecutionContext context) {

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

@@ -14,10 +14,7 @@
 
 import org.apiguardian.api.API;
 import org.junit.jupiter.api.extension.TestInstances;
-import org.junit.jupiter.engine.extension.ExtensionRegistrar;
 import org.junit.jupiter.engine.extension.ExtensionRegistry;
-import org.junit.jupiter.engine.extension.MutableExtensionRegistry;
-import org.junit.platform.engine.support.hierarchical.ThrowableCollector;
 
 /**
  * @since 5.0
@@ -26,12 +23,10 @@
 @API(status = INTERNAL, since = "5.0")
 public interface TestInstancesProvider {
 
-	default TestInstances getTestInstances(MutableExtensionRegistry extensionRegistry,
-			ThrowableCollector throwableCollector) {
-		return getTestInstances(extensionRegistry, extensionRegistry, throwableCollector);
+	default TestInstances getTestInstances(JupiterEngineExecutionContext context) {
+		return getTestInstances(context.getExtensionRegistry(), context);
 	}
 
-	TestInstances getTestInstances(ExtensionRegistry extensionRegistry, ExtensionRegistrar extensionRegistrar,
-			ThrowableCollector throwableCollector);
+	TestInstances getTestInstances(ExtensionRegistry extensionRegistry, JupiterEngineExecutionContext executionContext);
 
 }

jupiter-tests/src/test/java/org/junit/jupiter/engine/TestInstanceLifecycleTests.java

@@ -992,8 +992,7 @@ public void postProcessTestInstance(Object testInstance, ExtensionContext contex
 			trackLifecycle(context);
 			assertThat(context.getTestInstance()).isNotPresent();
 			assertNotNull(testInstance);
-			instanceMap.put(postProcessTestInstanceKey(context.getRequiredTestClass()),
-				DefaultTestInstances.of(testInstance));
+			instanceMap.put(postProcessTestInstanceKey(testInstance.getClass()), DefaultTestInstances.of(testInstance));
 		}
 
 		@Override

jupiter-tests/src/test/java/org/junit/jupiter/engine/extension/ParameterResolverTests.java

@@ -258,6 +258,8 @@ static class ConstructorInjectionTestCase {
 		void test() {
 			assertNotNull(this.outerTestInfo);
 			assertNotNull(this.outerCustomType);
+			assertEquals("test()", outerTestInfo.getDisplayName());
+			assertEquals(ConstructorInjectionTestCase.class, outerTestInfo.getTestClass().orElse(null));
 		}
 
 		@Nested
@@ -275,8 +277,13 @@ class NestedTestCase {
 			void test() {
 				assertNotNull(outerTestInfo);
 				assertNotNull(outerCustomType);
+				assertEquals("test()", outerTestInfo.getDisplayName());
+				assertEquals(NestedTestCase.class, outerTestInfo.getTestClass().orElse(null));
+
 				assertNotNull(this.innerTestInfo);
 				assertNotNull(this.innerCustomType);
+				assertEquals("test()", innerTestInfo.getDisplayName());
+				assertEquals(NestedTestCase.class, innerTestInfo.getTestClass().orElse(null));
 			}
 		}
 	}
@@ -296,6 +303,9 @@ static class AnnotatedParameterConstructorInjectionTestCase {
 		void test() {
 			assertNotNull(this.outerTestInfo);
 			assertNotNull(this.outerCustomType);
+			assertEquals("test()", outerTestInfo.getDisplayName());
+			assertEquals(AnnotatedParameterConstructorInjectionTestCase.class,
+				outerTestInfo.getTestClass().orElse(null));
 		}
 
 		@Nested
@@ -314,8 +324,15 @@ class AnnotatedConstructorParameterNestedTestCase {
 			void test() {
 				assertNotNull(outerTestInfo);
 				assertNotNull(outerCustomType);
+				assertEquals("test()", outerTestInfo.getDisplayName());
+				assertEquals(AnnotatedConstructorParameterNestedTestCase.class,
+					outerTestInfo.getTestClass().orElse(null));
+
 				assertNotNull(this.innerTestInfo);
 				assertNotNull(this.innerCustomType);
+				assertEquals("test()", innerTestInfo.getDisplayName());
+				assertEquals(AnnotatedConstructorParameterNestedTestCase.class,
+					innerTestInfo.getTestClass().orElse(null));
 			}
 		}
 	}

jupiter-tests/src/test/java/org/junit/jupiter/engine/extension/TestInstanceFactoryTests.java

@@ -276,7 +276,8 @@ void instanceFactoryOnTopLevelTestClass() {
 		// @formatter:off
 		assertThat(callSequence).containsExactly(
 			"FooInstanceFactory instantiated: ParentTestCase",
-				"parentTest"
+				"parentTest",
+			"close ParentTestCase"
 		);
 		// @formatter:on
 	}
@@ -305,8 +306,10 @@ void inheritedFactoryInTestClassHierarchy() {
 		assertThat(callSequence).containsExactly(
 			"FooInstanceFactory instantiated: InheritedFactoryTestCase",
 				"parentTest",
+			"close InheritedFactoryTestCase",
 			"FooInstanceFactory instantiated: InheritedFactoryTestCase",
-				"childTest"
+				"childTest",
+			"close InheritedFactoryTestCase"
 		);
 		// @formatter:on
 	}
@@ -324,17 +327,23 @@ void instanceFactoriesInNestedClassStructureAreInherited() {
 			// OuterTestCase
 			"FooInstanceFactory instantiated: OuterTestCase",
 				"outerTest",
+			"close OuterTestCase",
 
 			// InnerTestCase
 			"FooInstanceFactory instantiated: OuterTestCase",
 				"FooInstanceFactory instantiated: InnerTestCase",
 					"innerTest1",
+				"close InnerTestCase",
+			"close OuterTestCase",
 
 			// InnerInnerTestCase
 			"FooInstanceFactory instantiated: OuterTestCase",
 				"FooInstanceFactory instantiated: InnerTestCase",
 					"FooInstanceFactory instantiated: InnerInnerTestCase",
-						"innerTest2"
+						"innerTest2",
+					"close InnerInnerTestCase",
+				"close InnerTestCase",
+			"close OuterTestCase"
 		);
 		// @formatter:on
 	}
@@ -349,7 +358,8 @@ void instanceFactoryRegisteredViaTestInterface() {
 		// @formatter:off
 		assertThat(callSequence).containsExactly(
 			"FooInstanceFactory instantiated: FactoryFromInterfaceTestCase",
-				"test"
+				"test",
+			"close FactoryFromInterfaceTestCase"
 		);
 		// @formatter:on
 	}
@@ -386,7 +396,8 @@ void instanceFactoryWithPerClassLifecycle() {
 						"test1",
 					"@BeforeEach",
 						"test2",
-				"@AfterAll"
+				"@AfterAll",
+			"close PerClassLifecycleTestCase"
 		);
 		// @formatter:on
 	}
@@ -633,6 +644,10 @@ public Object createTestInstance(TestInstanceFactoryContext factoryContext, Exte
 			Class<?> testClass = factoryContext.getTestClass();
 			instantiated(getClass(), testClass);
 
+			extensionContext.getStore(ExtensionContext.Namespace.create(this)).put(new Object(),
+				(ExtensionContext.Store.CloseableResource) () -> callSequence.add(
+					"close " + testClass.getSimpleName()));
+
 			if (factoryContext.getOuterInstance().isPresent()) {
 				return ReflectionSupport.newInstance(testClass, factoryContext.getOuterInstance().get());
 			}