Provide test instance to container level extensions

Issue: #419

Author: sbrannen

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

@@ -28,8 +28,11 @@
 public final class ClassBasedContainerExtensionContext extends AbstractExtensionContext<ClassTestDescriptor>
 		implements ContainerExtensionContext {
 
+	private Object testInstance;
+
 	public ClassBasedContainerExtensionContext(ExtensionContext parent, EngineExecutionListener engineExecutionListener,
 			ClassTestDescriptor testDescriptor) {
+
 		super(parent, engineExecutionListener, testDescriptor);
 	}
 
@@ -43,10 +46,13 @@ public Optional<Class<?>> getTestClass() {
 		return Optional.of(getTestDescriptor().getTestClass());
 	}
 
+	void setTestInstance(Object testInstance) {
+		this.testInstance = testInstance;
+	}
+
 	@Override
 	public Optional<Object> getTestInstance() {
-		// TODO Determine how to access the test instance
-		return null;
+		return Optional.ofNullable(this.testInstance);
 	}
 
 	@Override

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

@@ -133,12 +133,21 @@ public JupiterEngineExecutionContext prepare(JupiterEngineExecutionContext conte
 		registerBeforeEachMethodAdapters(registry);
 		registerAfterEachMethodAdapters(registry);
 
-		ContainerExtensionContext containerExtensionContext = new ClassBasedContainerExtensionContext(
+		ClassBasedContainerExtensionContext containerExtensionContext = new ClassBasedContainerExtensionContext(
 			context.getExtensionContext(), context.getExecutionListener(), this);
 
+		// Reuse TestInstanceProvider for potential transparent instance caching.
+		TestInstanceProvider testInstanceProvider = testInstanceProvider(context, registry, containerExtensionContext);
+
+		// Eagerly load test instance for BeforeAllCallbacks, if necessary,
+		// and store the instance in the ContainerExtensionContext.
+		Object testInstance = (this.lifecycle == Lifecycle.PER_CLASS
+				? testInstanceProvider.getTestInstance(Optional.empty()) : null);
+		containerExtensionContext.setTestInstance(testInstance);
+
 		// @formatter:off
 		return context.extend()
-				.withTestInstanceProvider(testInstanceProvider(context, registry, containerExtensionContext))
+				.withTestInstanceProvider(testInstanceProvider)
 				.withExtensionRegistry(registry)
 				.withExtensionContext(containerExtensionContext)
 				.withThrowableCollector(new ThrowableCollector())
@@ -214,7 +223,7 @@ private void invokeBeforeAllMethods(JupiterEngineExecutionContext context) {
 		ExtensionRegistry registry = context.getExtensionRegistry();
 		ContainerExtensionContext extensionContext = (ContainerExtensionContext) context.getExtensionContext();
 		ThrowableCollector throwableCollector = context.getThrowableCollector();
-		Object testInstance = getTestInstanceForClassLevelCallbacks(context);
+		Object testInstance = extensionContext.getTestInstance().orElse(null);
 
 		for (Method method : this.beforeAllMethods) {
 			throwableCollector.execute(
@@ -229,17 +238,12 @@ private void invokeAfterAllMethods(JupiterEngineExecutionContext context) {
 		ExtensionRegistry registry = context.getExtensionRegistry();
 		ContainerExtensionContext extensionContext = (ContainerExtensionContext) context.getExtensionContext();
 		ThrowableCollector throwableCollector = context.getThrowableCollector();
-		Object testInstance = getTestInstanceForClassLevelCallbacks(context);
+		Object testInstance = extensionContext.getTestInstance().orElse(null);
 
 		this.afterAllMethods.forEach(method -> throwableCollector.execute(
 			() -> executableInvoker.invoke(method, testInstance, extensionContext, registry)));
 	}
 
-	private Object getTestInstanceForClassLevelCallbacks(JupiterEngineExecutionContext context) {
-		return this.lifecycle == Lifecycle.PER_CLASS
-				? context.getTestInstanceProvider().getTestInstance(Optional.empty()) : null;
-	}
-
 	private void invokeAfterAllCallbacks(JupiterEngineExecutionContext context) {
 		ExtensionRegistry registry = context.getExtensionRegistry();
 		ContainerExtensionContext extensionContext = (ContainerExtensionContext) context.getExtensionContext();

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

@@ -35,7 +35,9 @@ public final class MethodBasedTestExtensionContext extends AbstractExtensionCont
 
 	public MethodBasedTestExtensionContext(ExtensionContext parent, EngineExecutionListener engineExecutionListener,
 			MethodTestDescriptor testDescriptor, Object testInstance, ThrowableCollector throwableCollector) {
+
 		super(parent, engineExecutionListener, testDescriptor);
+
 		this.testInstance = testInstance;
 		this.throwableCollector = throwableCollector;
 	}

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

@@ -124,6 +124,8 @@ private void performAssertions(Class<?> testClass, int containers, int tests, in
 
 		ExecutionEventRecorder eventRecorder = executeTestsForClass(testClass);
 
+		// eventRecorder.eventStream().forEach(System.out::println);
+
 		// @formatter:off
 		assertAll(
 			() -> assertEquals(containers, eventRecorder.getContainerStartedCount(), "# containers started"),