Ensure context caching works properly during AOT runtime in the TCF

Prior to this commit, the AOT runtime support in the Spring TestContext
Framework (TCF) relied on the MergedContextConfiguration for a given
test class being the same as during the AOT processing phase. However,
this is not always the case. For example, Spring Boot "disables"
selected `ContextCustomizer` implementations during AOT runtime
execution.

See 0f325f98b5

To address that, this commit ensures that context caching works
properly during AOT runtime execution even if the
MergedContextConfiguration differs from what was produced during the
AOT processing phase. Specifically, this commit introduces
AotMergedContextConfiguration which is a MergedContextConfiguration
implementation based on an AOT-generated ApplicationContextInitializer.

AotMergedContextConfiguration wraps the MergedContextConfiguration
built during AOT runtime execution.

Interactions with the ContextCache are performed using the
AotMergedContextConfiguration; whereas, the ApplicationContext is
loaded using the original MergedContextConfiguration.

This commit also introduces a ContextCustomizerFactory that emulates
the ImportsContextCustomizerFactory in Spring Boot's testing support.
BasicSpringJupiterImportedConfigTests uses @Import to verify that the
context customizer works, and AotIntegrationTests has been updated to
execute BasicSpringJupiterImportedConfigTests after test classes whose
MergedContextConfiguration is identical during AOT runtime execution.
Without the fix in this commit, BasicSpringJupiterImportedConfigTests
would fail in AOT runtime mode since its ApplicationContext would be
pulled from the cache using an inappropriate cache key.

Closes gh-29289

spring-test/src/main/java/org/springframework/test/context/cache/AotMergedContextConfiguration.java

@@ -0,0 +1,112 @@
+/*
+ * 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.cache;
+
+import java.util.Collections;
+
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.core.style.ToStringCreator;
+import org.springframework.lang.Nullable;
+import org.springframework.test.context.CacheAwareContextLoaderDelegate;
+import org.springframework.test.context.MergedContextConfiguration;
+
+/**
+ * {@link MergedContextConfiguration} implementation based on an AOT-generated
+ * {@link ApplicationContextInitializer} that is used to load an AOT-optimized
+ * {@link org.springframework.context.ApplicationContext ApplicationContext}.
+ *
+ * <p>An {@code ApplicationContext} should not be loaded using the metadata in
+ * this {@code AotMergedContextConfiguration}. Rather the metadata from the
+ * {@linkplain #getOriginal() original} {@code MergedContextConfiguration} must
+ * be used.
+ *
+ * @author Sam Brannen
+ * @since 6.0
+ */
+class AotMergedContextConfiguration extends MergedContextConfiguration {
+
+	private static final long serialVersionUID = 1963364911008547843L;
+
+	private final Class<? extends ApplicationContextInitializer<?>> contextInitializerClass;
+
+	private final MergedContextConfiguration original;
+
+
+	AotMergedContextConfiguration(Class<?> testClass,
+			Class<? extends ApplicationContextInitializer<?>> contextInitializerClass,
+			MergedContextConfiguration original,
+			CacheAwareContextLoaderDelegate cacheAwareContextLoaderDelegate) {
+
+		super(testClass, null, null, Collections.singleton(contextInitializerClass), null,
+				original.getContextLoader(), cacheAwareContextLoaderDelegate, original.getParent());
+		this.contextInitializerClass = contextInitializerClass;
+		this.original = original;
+	}
+
+
+	/**
+	 * Get the original {@link MergedContextConfiguration} that this
+	 * {@code AotMergedContextConfiguration} was created for.
+	 */
+	MergedContextConfiguration getOriginal() {
+		return this.original;
+	}
+
+	@Override
+	public boolean equals(@Nullable Object other) {
+		if (this == other) {
+			return true;
+		}
+		if (other == null || other.getClass() != getClass()) {
+			return false;
+		}
+		AotMergedContextConfiguration that = (AotMergedContextConfiguration) other;
+		if (!this.contextInitializerClass.equals(that.contextInitializerClass)) {
+			return false;
+		}
+		if (!nullSafeClassName(getContextLoader()).equals(nullSafeClassName(that.getContextLoader()))) {
+			return false;
+		}
+		if (getParent() == null) {
+			if (that.getParent() != null) {
+				return false;
+			}
+		}
+		else if (!getParent().equals(that.getParent())) {
+			return false;
+		}
+		return true;
+	}
+
+	@Override
+	public int hashCode() {
+		int result = this.contextInitializerClass.hashCode();
+		result = 31 * result + nullSafeClassName(getContextLoader()).hashCode();
+		result = 31 * result + (getParent() != null ? getParent().hashCode() : 0);
+		return result;
+	}
+
+	@Override
+	public String toString() {
+		return new ToStringCreator(this)
+				.append("testClass", getTestClass().getName())
+				.append("contextInitializerClass", this.contextInitializerClass.getName())
+				.append("original", this.original)
+				.toString();
+	}
+
+}

spring-test/src/main/java/org/springframework/test/context/cache/DefaultCacheAwareContextLoaderDelegate.java

@@ -86,18 +86,19 @@ public class DefaultCacheAwareContextLoaderDelegate implements CacheAwareContext
 	@Override
 	public boolean isContextLoaded(MergedContextConfiguration mergedContextConfiguration) {
 		synchronized (this.contextCache) {
-			return this.contextCache.contains(mergedContextConfiguration);
+			return this.contextCache.contains(replaceIfNecessary(mergedContextConfiguration));
 		}
 	}
 
 	@Override
 	public ApplicationContext loadContext(MergedContextConfiguration mergedContextConfiguration) {
+		mergedContextConfiguration = replaceIfNecessary(mergedContextConfiguration);
 		synchronized (this.contextCache) {
 			ApplicationContext context = this.contextCache.get(mergedContextConfiguration);
 			if (context == null) {
 				try {
-					if (runningInAotMode(mergedContextConfiguration.getTestClass())) {
-						context = loadContextInAotMode(mergedContextConfiguration);
+					if (mergedContextConfiguration instanceof AotMergedContextConfiguration aotMergedConfig) {
+						context = loadContextInAotMode(aotMergedConfig);
 					}
 					else {
 						context = loadContextInternal(mergedContextConfiguration);
@@ -129,7 +130,7 @@ public class DefaultCacheAwareContextLoaderDelegate implements CacheAwareContext
 	@Override
 	public void closeContext(MergedContextConfiguration mergedContextConfiguration, @Nullable HierarchyMode hierarchyMode) {
 		synchronized (this.contextCache) {
-			this.contextCache.remove(mergedContextConfiguration, hierarchyMode);
+			this.contextCache.remove(replaceIfNecessary(mergedContextConfiguration), hierarchyMode);
 		}
 	}
 
@@ -163,23 +164,23 @@ public class DefaultCacheAwareContextLoaderDelegate implements CacheAwareContext
 		}
 	}
 
-	protected ApplicationContext loadContextInAotMode(MergedContextConfiguration mergedConfig) throws Exception {
-		Class<?> testClass = mergedConfig.getTestClass();
+	protected ApplicationContext loadContextInAotMode(AotMergedContextConfiguration aotMergedConfig) throws Exception {
+		Class<?> testClass = aotMergedConfig.getTestClass();
 		ApplicationContextInitializer<ConfigurableApplicationContext> contextInitializer =
 				this.aotTestContextInitializers.getContextInitializer(testClass);
 		Assert.state(contextInitializer != null,
 				() -> "Failed to load AOT ApplicationContextInitializer for test class [%s]"
 						.formatted(testClass.getName()));
-		ContextLoader contextLoader = getContextLoader(mergedConfig);
-		logger.info(LogMessage.format("Loading ApplicationContext in AOT mode for %s", mergedConfig));
+		ContextLoader contextLoader = getContextLoader(aotMergedConfig);
+		logger.info(LogMessage.format("Loading ApplicationContext in AOT mode for %s", aotMergedConfig.getOriginal()));
 		if (!((contextLoader instanceof AotContextLoader aotContextLoader) &&
-				(aotContextLoader.loadContextForAotRuntime(mergedConfig, contextInitializer)
+				(aotContextLoader.loadContextForAotRuntime(aotMergedConfig.getOriginal(), contextInitializer)
 						instanceof GenericApplicationContext gac))) {
 			throw new TestContextAotException("""
 					Cannot load ApplicationContext for AOT runtime for %s. The configured \
 					ContextLoader [%s] must be an AotContextLoader and must create a \
 					GenericApplicationContext."""
-						.formatted(mergedConfig, contextLoader.getClass().getName()));
+						.formatted(aotMergedConfig.getOriginal(), contextLoader.getClass().getName()));
 		}
 		gac.registerShutdownHook();
 		return gac;
@@ -195,10 +196,24 @@ public class DefaultCacheAwareContextLoaderDelegate implements CacheAwareContext
 	}
 
 	/**
-	 * Determine if we are running in AOT mode for the supplied test class.
+	 * If the test class associated with the supplied {@link MergedContextConfiguration}
+	 * has an AOT-optimized {@link ApplicationContext}, this method will create an
+	 * {@link AotMergedContextConfiguration} to replace the provided {@code MergedContextConfiguration}.
+	 * <p>Otherwise, this method simply returns the supplied {@code MergedContextConfiguration}
+	 * unmodified.
+	 * <p>This allows for transparent {@link org.springframework.test.context.cache.ContextCache ContextCache}
+	 * support for AOT-optimized application contexts.
 	 */
-	private boolean runningInAotMode(Class<?> testClass) {
-		return this.aotTestContextInitializers.isSupportedTestClass(testClass);
+	@SuppressWarnings("unchecked")
+	private MergedContextConfiguration replaceIfNecessary(MergedContextConfiguration mergedConfig) {
+		Class<?> testClass = mergedConfig.getTestClass();
+		if (this.aotTestContextInitializers.isSupportedTestClass(testClass)) {
+			Class<? extends ApplicationContextInitializer<?>> contextInitializerClass =
+					(Class<? extends ApplicationContextInitializer<?>>)
+							this.aotTestContextInitializers.getContextInitializer(testClass).getClass();
+			return new AotMergedContextConfiguration(testClass, contextInitializerClass, mergedConfig, this);
+		}
+		return mergedConfig;
 	}
 
 }