Introduce AotTestMappings and AotTestMappingsCodeGenerator

TestContextAotGenerator now uses AotTestMappingsCodeGenerator to
generate a AotTestMappings__Generated.java class which is loaded in
AotTestMappings via reflection in order to retrieve access to
ApplicationContextIntializers generated during AOT processing.

Furthermore, the processAheadOfTimeAndGenerateAotTestMappings() method
in TestContextAotGeneratorTests now performs a rather extensive test
including:

- emulating TestClassScanner to find test classes
- processing all test classes and generating ApplicationContextIntializers
- generating mappings for AotTestMappings
- compiling all generated code
- loading AotTestMappings
- using AotTestMappings to instantiate the generated ApplicationContextIntializer
- using the AotContextLoader API to load the AOT-optimized ApplicationContext
- asserting the behavior of the loaded ApplicationContext

See gh-28205
Closes gh-28204

spring-test/src/main/java/org/springframework/test/context/aot/AotTestMappings.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2002-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.test.context.aot;
+
+import java.lang.reflect.Method;
+import java.util.Map;
+import java.util.function.Supplier;
+
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
+import org.springframework.util.ReflectionUtils;
+
+/**
+ * {@code AotTestMappings} provides mappings from test classes to AOT-optimized
+ * context initializers.
+ *
+ * <p>If a test class is not {@linkplain #isSupportedTestClass(Class) supported} in
+ * AOT mode, {@link #getContextInitializer(Class)} will return {@code null}.
+ *
+ * <p>Reflectively accesses {@link #GENERATED_MAPPINGS_CLASS_NAME} generated by
+ * the {@link TestContextAotGenerator} to retrieve the mappings generated during
+ * AOT processing.
+ *
+ * @author Sam Brannen
+ * @author Stephane Nicoll
+ * @since 6.0
+ */
+public class AotTestMappings {
+
+	// TODO Add support in ClassNameGenerator for supplying a predefined class name.
+	// There is a similar issue in Spring Boot where code relies on a generated name.
+	// Ideally we would generate a class named: org.springframework.test.context.aot.GeneratedAotTestMappings
+	static final String GENERATED_MAPPINGS_CLASS_NAME = AotTestMappings.class.getName() + "__Generated";
+
+	static final String GENERATED_MAPPINGS_METHOD_NAME = "getContextInitializers";
+
+	private final Map<String, Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>>> contextInitializers;
+
+
+	public AotTestMappings() {
+		this(GENERATED_MAPPINGS_CLASS_NAME);
+	}
+
+	AotTestMappings(String initializerClassName) {
+		this(loadContextInitializersMap(initializerClassName));
+	}
+
+	AotTestMappings(Map<String, Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>>> contextInitializers) {
+		this.contextInitializers = contextInitializers;
+	}
+
+
+	/**
+	 * Determine if the specified test class has an AOT-optimized application context
+	 * initializer.
+	 * <p>If this method returns {@code true}, {@link #getContextInitializer(Class)}
+	 * should not return {@code null}.
+	 */
+	public boolean isSupportedTestClass(Class<?> testClass) {
+		return this.contextInitializers.containsKey(testClass.getName());
+	}
+
+	/**
+	 * Get the AOT {@link ApplicationContextInitializer} for the specified test class.
+	 * @return the AOT context initializer, or {@code null} if there is no AOT context
+	 * initializer for the specified test class
+	 * @see #isSupportedTestClass(Class)
+	 */
+	public ApplicationContextInitializer<ConfigurableApplicationContext> getContextInitializer(Class<?> testClass) {
+		Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>> supplier =
+				this.contextInitializers.get(testClass.getName());
+		return (supplier != null ? supplier.get() : null);
+	}
+
+
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	private static Map<String, Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>>>
+			loadContextInitializersMap(String className) {
+
+		String methodName = GENERATED_MAPPINGS_METHOD_NAME;
+
+		try {
+			Class<?> clazz = ClassUtils.forName(className, null);
+			Method method = ReflectionUtils.findMethod(clazz, methodName);
+			Assert.state(method != null, () -> "No %s() method found in %s".formatted(methodName, clazz.getName()));
+			return (Map<String, Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>>>)
+					ReflectionUtils.invokeMethod(method, null);
+		}
+		catch (IllegalStateException ex) {
+			throw ex;
+		}
+		catch (Exception ex) {
+			throw new IllegalStateException("Failed to invoke %s() method in %s".formatted(methodName, className), ex);
+		}
+	}
+
+}

spring-test/src/main/java/org/springframework/test/context/aot/AotTestMappingsCodeGenerator.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2002-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.test.context.aot;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Supplier;
+
+import javax.lang.model.element.Modifier;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+import org.springframework.aot.generate.GeneratedClass;
+import org.springframework.aot.generate.GeneratedClasses;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.core.log.LogMessage;
+import org.springframework.javapoet.ClassName;
+import org.springframework.javapoet.CodeBlock;
+import org.springframework.javapoet.MethodSpec;
+import org.springframework.javapoet.ParameterizedTypeName;
+import org.springframework.javapoet.TypeName;
+import org.springframework.javapoet.TypeSpec;
+import org.springframework.javapoet.WildcardTypeName;
+import org.springframework.util.MultiValueMap;
+
+/**
+ * Internal code generator for mappings used by {@link AotTestMappings}.
+ *
+ * @author Sam Brannen
+ * @since 6.0
+ */
+class AotTestMappingsCodeGenerator {
+
+	private static final Log logger = LogFactory.getLog(AotTestMappingsCodeGenerator.class);
+
+	private static final ParameterizedTypeName CONTEXT_INITIALIZER = ParameterizedTypeName.get(
+			ClassName.get(ApplicationContextInitializer.class),
+			WildcardTypeName.subtypeOf(ConfigurableApplicationContext.class));
+
+	private static final ParameterizedTypeName CONTEXT_INITIALIZER_SUPPLIER = ParameterizedTypeName
+			.get(ClassName.get(Supplier.class), CONTEXT_INITIALIZER);
+
+	// Map<String, Supplier<ApplicationContextInitializer<? extends ConfigurableApplicationContext>>>
+	private static final TypeName CONTEXT_SUPPLIER_MAP = ParameterizedTypeName
+			.get(ClassName.get(Map.class), ClassName.get(String.class), CONTEXT_INITIALIZER_SUPPLIER);
+
+
+	private final MultiValueMap<ClassName, Class<?>> initializerClassMappings;
+
+	private final GeneratedClass generatedClass;
+
+
+	AotTestMappingsCodeGenerator(MultiValueMap<ClassName, Class<?>> initializerClassMappings,
+			GeneratedClasses generatedClasses) {
+
+		this.initializerClassMappings = initializerClassMappings;
+		this.generatedClass = generatedClasses.addForFeature("Generated", this::generateType);
+	}
+
+
+	GeneratedClass getGeneratedClass() {
+		return this.generatedClass;
+	}
+
+	private void generateType(TypeSpec.Builder type) {
+		logger.debug(LogMessage.format("Generating AOT test mappings in %s",
+				this.generatedClass.getName().reflectionName()));
+		type.addJavadoc("Generated mappings for {@link $T}.", AotTestMappings.class);
+		type.addModifiers(Modifier.PUBLIC);
+		type.addMethod(generateMappingMethod());
+	}
+
+	private MethodSpec generateMappingMethod() {
+		MethodSpec.Builder method = MethodSpec.methodBuilder(AotTestMappings.GENERATED_MAPPINGS_METHOD_NAME);
+		method.addModifiers(Modifier.PUBLIC, Modifier.STATIC);
+		method.returns(CONTEXT_SUPPLIER_MAP);
+		method.addCode(generateMappingCode());
+		return method.build();
+	}
+
+	private CodeBlock generateMappingCode() {
+		CodeBlock.Builder code = CodeBlock.builder();
+		code.addStatement("$T map = new $T<>()", CONTEXT_SUPPLIER_MAP, HashMap.class);
+		this.initializerClassMappings.forEach((className, testClasses) -> {
+			List<String> testClassNames = testClasses.stream().map(Class::getName).toList();
+			logger.debug(LogMessage.format(
+					"Generating mapping from AOT context initializer [%s] to test classes %s",
+					className.reflectionName(), testClassNames));
+			testClassNames.forEach(testClassName ->
+				code.addStatement("map.put($S, () -> new $T())", testClassName, className));
+		});
+		code.addStatement("return map");
+		return code.build();
+	}
+
+}

spring-test/src/main/java/org/springframework/test/context/aot/TestContextAotGenerator.java

@@ -24,13 +24,17 @@ import org.apache.commons.logging.LogFactory;
 
 import org.springframework.aot.generate.ClassNameGenerator;
 import org.springframework.aot.generate.DefaultGenerationContext;
+import org.springframework.aot.generate.GeneratedClasses;
 import org.springframework.aot.generate.GeneratedFiles;
 import org.springframework.aot.generate.GenerationContext;
+import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.RuntimeHints;
+import org.springframework.aot.hint.TypeReference;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextInitializer;
 import org.springframework.context.aot.ApplicationContextAotGenerator;
 import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.core.log.LogMessage;
 import org.springframework.javapoet.ClassName;
 import org.springframework.test.context.BootstrapUtils;
 import org.springframework.test.context.ContextLoader;
@@ -51,7 +55,7 @@ import org.springframework.util.MultiValueMap;
  */
 class TestContextAotGenerator {
 
-	private static final Log logger = LogFactory.getLog(TestClassScanner.class);
+	private static final Log logger = LogFactory.getLog(TestContextAotGenerator.class);
 
 	private final ApplicationContextAotGenerator aotGenerator = new ApplicationContextAotGenerator();
 
@@ -97,30 +101,33 @@ class TestContextAotGenerator {
 	 * @throws TestContextAotException if an error occurs during AOT processing
 	 */
 	public void processAheadOfTime(Stream<Class<?>> testClasses) throws TestContextAotException {
-		MultiValueMap<MergedContextConfiguration, Class<?>> map = new LinkedMultiValueMap<>();
-		testClasses.forEach(testClass -> map.add(buildMergedContextConfiguration(testClass), testClass));
-
-		map.forEach((mergedConfig, classes) -> {
-			// System.err.println(mergedConfig + " -> " + classes);
-			if (logger.isDebugEnabled()) {
-				logger.debug("Generating AOT artifacts for test classes [%s]"
-						.formatted(classes.stream().map(Class::getCanonicalName).toList()));
-			}
+		MultiValueMap<MergedContextConfiguration, Class<?>> mergedConfigMappings = new LinkedMultiValueMap<>();
+		testClasses.forEach(testClass -> mergedConfigMappings.add(buildMergedContextConfiguration(testClass), testClass));
+		processAheadOfTime(mergedConfigMappings);
+	}
+
+	private void processAheadOfTime(MultiValueMap<MergedContextConfiguration, Class<?>> mergedConfigMappings) {
+		MultiValueMap<ClassName, Class<?>> initializerClassMappings = new LinkedMultiValueMap<>();
+		mergedConfigMappings.forEach((mergedConfig, testClasses) -> {
+			logger.debug(LogMessage.format("Generating AOT artifacts for test classes %s",
+					testClasses.stream().map(Class::getName).toList()));
 			try {
 				// Use first test class discovered for a given unique MergedContextConfiguration.
-				Class<?> testClass = classes.get(0);
+				Class<?> testClass = testClasses.get(0);
 				DefaultGenerationContext generationContext = createGenerationContext(testClass);
-				ClassName className = processAheadOfTime(mergedConfig, generationContext);
-				// TODO Store ClassName in a map analogous to TestContextAotProcessor in Spring Native.
+				ClassName initializer = processAheadOfTime(mergedConfig, generationContext);
+				Assert.state(!initializerClassMappings.containsKey(initializer),
+						() -> "ClassName [%s] already encountered".formatted(initializer.reflectionName()));
+				initializerClassMappings.addAll(initializer, testClasses);
 				generationContext.writeGeneratedContent();
 			}
 			catch (Exception ex) {
-				if (logger.isWarnEnabled()) {
-					logger.warn("Failed to generate AOT artifacts for test classes [%s]"
-							.formatted(classes.stream().map(Class::getCanonicalName).toList()), ex);
-				}
+				logger.warn(LogMessage.format("Failed to generate AOT artifacts for test classes [%s]",
+						testClasses.stream().map(Class::getName).toList()), ex);
 			}
 		});
+
+		generateAotTestMappings(initializerClassMappings);
 	}
 
 	/**
@@ -143,7 +150,7 @@ class TestContextAotGenerator {
 		}
 		catch (Throwable ex) {
 			throw new TestContextAotException("Failed to process test class [%s] for AOT"
-					.formatted(mergedConfig.getTestClass().getCanonicalName()), ex);
+					.formatted(mergedConfig.getTestClass().getName()), ex);
 		}
 	}
 
@@ -154,7 +161,7 @@ class TestContextAotGenerator {
 	 * create {@link GenericApplicationContext GenericApplicationContexts}.
 	 * @throws TestContextAotException if an error occurs while loading the application
 	 * context or if one of the prerequisites is not met
-	 * @see SmartContextLoader#loadContextForAotProcessing(MergedContextConfiguration)
+	 * @see AotContextLoader#loadContextForAotProcessing(MergedContextConfiguration)
 	 */
 	private GenericApplicationContext loadContextForAotProcessing(
 			MergedContextConfiguration mergedConfig) throws TestContextAotException {
@@ -164,7 +171,7 @@ class TestContextAotGenerator {
 		Assert.notNull(contextLoader, """
 				Cannot load an ApplicationContext with a NULL 'contextLoader'. \
 				Consider annotating test class [%s] with @ContextConfiguration or \
-				@ContextHierarchy.""".formatted(testClass.getCanonicalName()));
+				@ContextHierarchy.""".formatted(testClass.getName()));
 
 		if (contextLoader instanceof AotContextLoader aotContextLoader) {
 			try {
@@ -176,13 +183,13 @@ class TestContextAotGenerator {
 			catch (Exception ex) {
 				throw new TestContextAotException(
 						"Failed to load ApplicationContext for AOT processing for test class [%s]"
-							.formatted(testClass.getCanonicalName()), ex);
+							.formatted(testClass.getName()), ex);
 			}
 		}
 		throw new TestContextAotException("""
 				Cannot generate AOT artifacts for test class [%s]. The configured \
 				ContextLoader [%s] must be an AotContextLoader and must create a \
-				GenericApplicationContext.""".formatted(testClass.getCanonicalName(),
+				GenericApplicationContext.""".formatted(testClass.getName(),
 					contextLoader.getClass().getName()));
 	}
 
@@ -203,4 +210,18 @@ class TestContextAotGenerator {
 		return "TestContext%03d_".formatted(this.sequence.incrementAndGet());
 	}
 
+	private void generateAotTestMappings(MultiValueMap<ClassName, Class<?>> initializerClassMappings) {
+		ClassNameGenerator classNameGenerator = new ClassNameGenerator(AotTestMappings.class);
+		DefaultGenerationContext generationContext =
+				new DefaultGenerationContext(classNameGenerator, this.generatedFiles, this.runtimeHints);
+		GeneratedClasses generatedClasses = generationContext.getGeneratedClasses();
+
+		AotTestMappingsCodeGenerator codeGenerator =
+				new AotTestMappingsCodeGenerator(initializerClassMappings, generatedClasses);
+		generationContext.writeGeneratedContent();
+		String className = codeGenerator.getGeneratedClass().getName().reflectionName();
+		this.runtimeHints.reflection().registerType(TypeReference.of(className),
+				builder -> builder.withMembers(MemberCategory.INVOKE_PUBLIC_METHODS));
+	}
+
 }

spring-test/src/test/java/org/springframework/test/context/aot/AbstractAotTests.java

@@ -28,6 +28,8 @@ import java.util.stream.Stream;
 abstract class AbstractAotTests {
 
 	static final String[] expectedSourceFilesForBasicSpringTests = {
+			// Global
+			"org/springframework/test/context/aot/AotTestMappings__Generated.java",
 			// BasicSpringJupiterSharedConfigTests
 			"org/springframework/context/event/DefaultEventListenerFactory__TestContext001_BeanDefinitions.java",
 			"org/springframework/context/event/EventListenerMethodProcessor__TestContext001_BeanDefinitions.java",

spring-test/src/test/java/org/springframework/test/context/aot/AotSmokeTests.java

@@ -49,9 +49,11 @@ class AotSmokeTests extends AbstractAotTests {
 		List<String> sourceFiles = generatedFiles.getGeneratedFiles(Kind.SOURCE).keySet().stream().toList();
 		assertThat(sourceFiles).containsExactlyInAnyOrder(expectedSourceFilesForBasicSpringTests);
 
-		TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled -> {
-			// just make sure compilation completes without errors
-		});
+		TestCompiler.forSystem().withFiles(generatedFiles)
+			// .printFiles(System.out)
+			.compile(compiled -> {
+				// just make sure compilation completes without errors
+			});
 	}
 
 }

spring-test/src/test/java/org/springframework/test/context/aot/TestContextAotGeneratorTests.java

@@ -19,18 +19,20 @@ package org.springframework.test.context.aot;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
-import java.util.stream.Stream;
 
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aot.generate.DefaultGenerationContext;
 import org.springframework.aot.generate.GeneratedFiles.Kind;
 import org.springframework.aot.generate.InMemoryGeneratedFiles;
+import org.springframework.aot.hint.MemberCategory;
+import org.springframework.aot.hint.ReflectionHints;
+import org.springframework.aot.hint.TypeReference;
 import org.springframework.aot.test.generator.compile.CompileWithTargetClassAccess;
 import org.springframework.aot.test.generator.compile.TestCompiler;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextInitializer;
-import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.javapoet.ClassName;
 import org.springframework.test.context.MergedContextConfiguration;
 import org.springframework.test.context.aot.samples.basic.BasicSpringJupiterSharedConfigTests;
@@ -48,6 +50,7 @@ import org.springframework.test.web.servlet.MockMvc;
 import org.springframework.util.function.ThrowingConsumer;
 import org.springframework.web.context.WebApplicationContext;
 
+import static java.util.Comparator.comparing;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.content;
@@ -67,8 +70,8 @@ class TestContextAotGeneratorTests extends AbstractAotTests {
 	 * @see AotSmokeTests#scanClassPathThenGenerateSourceFilesAndCompileThem()
 	 */
 	@Test
-	void generate() {
-		Stream<Class<?>> testClasses = Stream.of(
+	void processAheadOfTimeAndGenerateAotTestMappings() {
+		Set<Class<?>> testClasses = Set.of(
 				BasicSpringJupiterSharedConfigTests.class,
 				BasicSpringJupiterTests.class,
 				BasicSpringJupiterTests.NestedTests.class,
@@ -78,14 +81,28 @@ class TestContextAotGeneratorTests extends AbstractAotTests {
 		InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
 		TestContextAotGenerator generator = new TestContextAotGenerator(generatedFiles);
 
-		generator.processAheadOfTime(testClasses);
+		generator.processAheadOfTime(testClasses.stream().sorted(comparing(Class::getName)));
+
+		ReflectionHints reflectionHints = generator.getRuntimeHints().reflection();
+		assertThat(reflectionHints.getTypeHint(TypeReference.of(AotTestMappings.GENERATED_MAPPINGS_CLASS_NAME)))
+				.satisfies(typeHint ->
+						assertThat(typeHint.getMemberCategories()).containsExactly(MemberCategory.INVOKE_PUBLIC_METHODS));
 
 		List<String> sourceFiles = generatedFiles.getGeneratedFiles(Kind.SOURCE).keySet().stream().toList();
 		assertThat(sourceFiles).containsExactlyInAnyOrder(expectedSourceFilesForBasicSpringTests);
 
-		TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled -> {
-			// just make sure compilation completes without errors
-		});
+		TestCompiler.forSystem().withFiles(generatedFiles).compile(ThrowingConsumer.of(compiled -> {
+			AotTestMappings aotTestMappings = new AotTestMappings();
+			for (Class<?> testClass : testClasses) {
+				MergedContextConfiguration mergedConfig = generator.buildMergedContextConfiguration(testClass);
+				ApplicationContextInitializer<ConfigurableApplicationContext> contextInitializer =
+						aotTestMappings.getContextInitializer(testClass);
+				assertThat(contextInitializer).isNotNull();
+				ApplicationContext context = ((AotContextLoader) mergedConfig.getContextLoader())
+						.loadContextForAotRuntime(mergedConfig, contextInitializer);
+				assertContextForBasicTests(context);
+			}
+		}));
 	}
 
 	@Test
@@ -99,13 +116,14 @@ class TestContextAotGeneratorTests extends AbstractAotTests {
 				BasicSpringTestNGTests.class,
 				BasicSpringVintageTests.class);
 
-		processAheadOfTime(testClasses, context -> {
-			assertThat(context.getEnvironment().getProperty("test.engine"))
-				.as("Environment").isNotNull();
+		processAheadOfTime(testClasses, this::assertContextForBasicTests);
+	}
 
-			MessageService messageService = context.getBean(MessageService.class);
-			assertThat(messageService.generateMessage()).isEqualTo("Hello, AOT!");
-		});
+	private void assertContextForBasicTests(ApplicationContext context) {
+		assertThat(context.getEnvironment().getProperty("test.engine")).as("Environment").isNotNull();
+
+		MessageService messageService = context.getBean(MessageService.class);
+		assertThat(messageService.generateMessage()).isEqualTo("Hello, AOT!");
 	}
 
 	@Test
@@ -151,16 +169,16 @@ class TestContextAotGeneratorTests extends AbstractAotTests {
 		InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
 		TestContextAotGenerator generator = new TestContextAotGenerator(generatedFiles);
 		List<Mapping> mappings = processAheadOfTime(generator, testClasses);
-		TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled -> {
-			mappings.forEach(mapping -> {
+		TestCompiler.forSystem().withFiles(generatedFiles).compile(ThrowingConsumer.of(compiled -> {
+			for (Mapping mapping : mappings) {
 				MergedContextConfiguration mergedConfig = mapping.mergedConfig();
-				ApplicationContextInitializer<GenericApplicationContext> contextInitializer =
+				ApplicationContextInitializer<ConfigurableApplicationContext> contextInitializer =
 						compiled.getInstance(ApplicationContextInitializer.class, mapping.className().reflectionName());
-				AotRuntimeContextLoader aotRuntimeContextLoader = new AotRuntimeContextLoader();
-				GenericApplicationContext context = aotRuntimeContextLoader.loadContext(mergedConfig, contextInitializer);
+				ApplicationContext context = ((AotContextLoader) mergedConfig.getContextLoader())
+						.loadContextForAotRuntime(mergedConfig, contextInitializer);
 				result.accept(context);
-			});
-		});
+			}
+		}));
 	}
 
 	private List<Mapping> processAheadOfTime(TestContextAotGenerator generator, Set<Class<?>> testClasses) {