Introduce AotContextLoader SPI in the TestContext framework

This commit introduces a dedicated AotContextLoader extension of
SmartContextLoader and reverts the recent changes to the existing
contract for SmartContextLoader. This allows existing
SmartContextLoader implementations to continue working unaffected by
AOT support unless they opt into AOT support by explicitly implementing
the new AotContextLoader contract.

In addition, existing SmartContextLoader implementations in the
spring-test module have been modified to implement AotContextLoader
instead of SmartContextLoader. This allows the core framework to
provide built-in AOT support in the TestContext framework, and it also
allows third-party extensions to built-in SmartContextLoaders to
participate in AOT processing and run-time support transparently (or at
least as transparent as possible).

Closes gh-28906

spring-test/src/main/java/org/springframework/test/context/SmartContextLoader.java

@@ -137,32 +137,6 @@ public interface SmartContextLoader extends ContextLoader {
 	 */
 	ApplicationContext loadContext(MergedContextConfiguration mergedConfig) throws Exception;
 
-	/**
-	 * Load a new {@linkplain ApplicationContext} based on the supplied
-	 * {@link MergedContextConfiguration}, configure the context, and return the
-	 * context.
-	 * <p>In contrast to {@link #loadContext(MergedContextConfiguration)}, this
-	 * method must <strong>not</strong>
-	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#refresh()
-	 * refresh} the {@code ApplicationContext} or
-	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#registerShutdownHook()
-	 * register a JVM shutdown hook} for it. Otherwise, this method should implement
-	 * behavior identical to {@link #loadContext(MergedContextConfiguration)}.
-	 * <p>The default implementation throws an {@link UnsupportedOperationException}.
-	 * Concrete implementations must therefore override this method in order to
-	 * support AOT (ahead of time) processing.
-	 * @param mergedConfig the merged context configuration to use to load the
-	 * application context
-	 * @return a new application context
-	 * @throws Exception if context loading failed
-	 * @since 6.0
-	 */
-	default ApplicationContext loadContextForAotProcessing(MergedContextConfiguration mergedConfig) throws Exception {
-		throw new UnsupportedOperationException(
-				"%s does not support loadContextForAotProcessing(MergedContextConfiguration)"
-					.formatted(getClass().getName()));
-	}
-
 	/**
 	 * {@code SmartContextLoader} does not support deprecated {@link ContextLoader} methods.
 	 * Call {@link #processContextConfiguration(ContextConfigurationAttributes)} instead.

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

@@ -0,0 +1,77 @@
+/*
+ * 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 org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.test.context.MergedContextConfiguration;
+import org.springframework.test.context.SmartContextLoader;
+
+/**
+ * Strategy interface for loading an {@link ApplicationContext} for build-time
+ * {@linkplain #loadContextForAotProcessing AOT processing} as well as run-time
+ * {@linkplain #loadContextForAotRuntime AOT execution} for an integration test
+ * managed by the Spring TestContext Framework.
+ *
+ * <p>As of Spring Framework 6.0, AOT infrastructure requires that an {@code AotContextLoader}
+ * create a {@link org.springframework.context.support.GenericApplicationContext
+ * GenericApplicationContext} for both build-time processing and run-time execution.
+ *
+ * @author Sam Brannen
+ * @since 6.0
+ */
+public interface AotContextLoader extends SmartContextLoader {
+
+	/**
+	 * Load a new {@link ApplicationContext} for AOT build-time processing based
+	 * on the supplied {@link MergedContextConfiguration}, configure the context,
+	 * and return the context.
+	 * <p>In contrast to {@link #loadContext(MergedContextConfiguration)}, this
+	 * method must <strong>not</strong>
+	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#refresh()
+	 * refresh} the {@code ApplicationContext} or
+	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#registerShutdownHook()
+	 * register a JVM shutdown hook} for it. Otherwise, this method should implement
+	 * behavior identical to {@link #loadContext(MergedContextConfiguration)}.
+	 * @param mergedConfig the merged context configuration to use to load the
+	 * application context
+	 * @return a new {@code GenericApplicationContext}
+	 * @throws Exception if context loading failed
+	 * @see #loadContextForAotRuntime(MergedContextConfiguration, ApplicationContextInitializer)
+	 */
+	ApplicationContext loadContextForAotProcessing(MergedContextConfiguration mergedConfig) throws Exception;
+
+	/**
+	 * Load a new {@link ApplicationContext} for AOT run-time execution based on
+	 * the supplied {@link MergedContextConfiguration} and
+	 * {@link ApplicationContextInitializer}.
+	 * <p>This method must instantiate, initialize, and
+	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#refresh()
+	 * refresh} the {@code ApplicationContext}.
+	 * @param mergedConfig the merged context configuration to use to load the
+	 * application context
+	 * @param initializer the {@code ApplicationContextInitializer} that should
+	 * be applied to the context in order to recreate bean definitions
+	 * @return a new {@code GenericApplicationContext}
+	 * @throws Exception if context loading failed
+	 * @see #loadContextForAotProcessing(MergedContextConfiguration)
+	 */
+	ApplicationContext loadContextForAotRuntime(MergedContextConfiguration mergedConfig,
+			ApplicationContextInitializer<ConfigurableApplicationContext> initializer) throws Exception;
+
+}

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

@@ -166,9 +166,9 @@ class TestContextAotGenerator {
 				Consider annotating test class [%s] with @ContextConfiguration or \
 				@ContextHierarchy.""".formatted(testClass.getCanonicalName()));
 
-		if (contextLoader instanceof SmartContextLoader smartContextLoader) {
+		if (contextLoader instanceof AotContextLoader aotContextLoader) {
 			try {
-				ApplicationContext context = smartContextLoader.loadContextForAotProcessing(mergedConfig);
+				ApplicationContext context = aotContextLoader.loadContextForAotProcessing(mergedConfig);
 				if (context instanceof GenericApplicationContext gac) {
 					return gac;
 				}
@@ -181,7 +181,7 @@ class TestContextAotGenerator {
 		}
 		throw new TestContextAotException("""
 				Cannot generate AOT artifacts for test class [%s]. The configured \
-				ContextLoader [%s] must be a SmartContextLoader and must create a \
+				ContextLoader [%s] must be an AotContextLoader and must create a \
 				GenericApplicationContext.""".formatted(testClass.getCanonicalName(),
 					contextLoader.getClass().getName()));
 	}

spring-test/src/main/java/org/springframework/test/context/support/AbstractDelegatingSmartContextLoader.java

@@ -20,11 +20,14 @@ import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
 import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.ContextConfigurationAttributes;
 import org.springframework.test.context.ContextLoader;
 import org.springframework.test.context.MergedContextConfiguration;
 import org.springframework.test.context.SmartContextLoader;
+import org.springframework.test.context.aot.AotContextLoader;
 import org.springframework.util.Assert;
 
 /**
@@ -68,7 +71,7 @@ import org.springframework.util.Assert;
  * @since 3.2
  * @see SmartContextLoader
  */
-public abstract class AbstractDelegatingSmartContextLoader implements SmartContextLoader {
+public abstract class AbstractDelegatingSmartContextLoader implements AotContextLoader {
 
 	private static final Log logger = LogFactory.getLog(AbstractDelegatingSmartContextLoader.class);
 
@@ -224,10 +227,16 @@ public abstract class AbstractDelegatingSmartContextLoader implements SmartConte
 	 * @since 6.0
 	 */
 	@Override
-	public ApplicationContext loadContextForAotProcessing(MergedContextConfiguration mergedConfig) throws Exception {
+	public final ApplicationContext loadContextForAotProcessing(MergedContextConfiguration mergedConfig) throws Exception {
 		return loadContext(mergedConfig, false);
 	}
 
+	@Override
+	public final ApplicationContext loadContextForAotRuntime(MergedContextConfiguration mergedConfig,
+			ApplicationContextInitializer<ConfigurableApplicationContext> initializer) throws Exception {
+		return getAotContextLoader(mergedConfig, "load").loadContextForAotRuntime(mergedConfig, initializer);
+	}
+
 	/**
 	 * Delegates to an appropriate candidate {@code SmartContextLoader} to load
 	 * an {@link ApplicationContext}.
@@ -240,12 +249,7 @@ public abstract class AbstractDelegatingSmartContextLoader implements SmartConte
 	 * {@code ApplicationContext} from the supplied merged context configuration
 	 */
 	private ApplicationContext loadContext(MergedContextConfiguration mergedConfig, boolean refresh) throws Exception {
-		Assert.notNull(mergedConfig, "MergedContextConfiguration must not be null");
-
-		Assert.state(!(mergedConfig.hasLocations() && mergedConfig.hasClasses()), () -> String.format(
-				"Neither %s nor %s supports loading an ApplicationContext from %s: " +
-				"declare either 'locations' or 'classes' but not both.", name(getXmlLoader()),
-				name(getAnnotationConfigLoader()), mergedConfig));
+		assertPreconditions(mergedConfig, "load");
 
 		SmartContextLoader[] candidates = {getXmlLoader(), getAnnotationConfigLoader()};
 		for (SmartContextLoader loader : candidates) {
@@ -265,7 +269,7 @@ public abstract class AbstractDelegatingSmartContextLoader implements SmartConte
 
 		// else...
 		throw new IllegalStateException(String.format(
-				"Neither %s nor %s was able to load an ApplicationContext from %s.",
+				"Neither %s nor %s is able to load an ApplicationContext for %s.",
 				name(getXmlLoader()), name(getAnnotationConfigLoader()), mergedConfig));
 	}
 
@@ -282,9 +286,53 @@ public abstract class AbstractDelegatingSmartContextLoader implements SmartConte
 			throws Exception {
 
 		if (logger.isDebugEnabled()) {
-			logger.debug(String.format("Delegating to %s to load context from %s.", name(loader), mergedConfig));
+			logger.debug(String.format("Delegating to %s to load context for %s.", name(loader), mergedConfig));
+		}
+		if (refresh) {
+			return loader.loadContext(mergedConfig);
+		}
+		else {
+			if (loader instanceof AotContextLoader aotContextLoader) {
+				return aotContextLoader.loadContextForAotProcessing(mergedConfig);
+			}
+			throw new IllegalStateException(
+					"%s must implement AotContextLoader to load the ApplicationContext for %s."
+						.formatted(name(loader), mergedConfig));
 		}
-		return (refresh ? loader.loadContext(mergedConfig) : loader.loadContextForAotProcessing(mergedConfig));
+	}
+
+	private AotContextLoader getAotContextLoader(MergedContextConfiguration mergedConfig, String action) {
+		assertPreconditions(mergedConfig, action);
+
+		SmartContextLoader[] candidates = {getXmlLoader(), getAnnotationConfigLoader()};
+		for (SmartContextLoader loader : candidates) {
+			// Determine if each loader can load a context from the mergedConfig. If it
+			// can, let it; otherwise, keep iterating.
+			if (loader instanceof AotContextLoader aotContextLoader && supports(loader, mergedConfig)) {
+				return aotContextLoader;
+			}
+		}
+
+		// If neither of the candidates supports the mergedConfig based on resources but
+		// ACIs or customizers were declared, then delegate to the annotation config
+		// loader.
+		if (getAnnotationConfigLoader() instanceof AotContextLoader aotContextLoader &&
+				(!mergedConfig.getContextInitializerClasses().isEmpty() ||
+						!mergedConfig.getContextCustomizers().isEmpty())) {
+			return aotContextLoader;
+		}
+
+		throw new IllegalStateException(
+				"Neither %s nor %s is able to %s the ApplicationContext for %s.".formatted(
+					name(getXmlLoader()), name(getAnnotationConfigLoader()), action, mergedConfig));
+	}
+
+	private void assertPreconditions(MergedContextConfiguration mergedConfig, String action) {
+		Assert.notNull(mergedConfig, "MergedContextConfiguration must not be null");
+		Assert.state(!(mergedConfig.hasLocations() && mergedConfig.hasClasses()), () -> """
+				Neither %s nor %s is able to %s an ApplicationContext for %s: \
+				declare either 'locations' or 'classes' but not both.""".formatted(
+						name(getXmlLoader()), name(getAnnotationConfigLoader()), action, mergedConfig));
 	}
 
 	private boolean supports(SmartContextLoader loader, MergedContextConfiguration mergedConfig) {

spring-test/src/main/java/org/springframework/test/context/support/AbstractGenericContextLoader.java

@@ -22,10 +22,12 @@ import org.apache.commons.logging.LogFactory;
 import org.springframework.beans.factory.support.BeanDefinitionReader;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextInitializer;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.annotation.AnnotationConfigUtils;
 import org.springframework.context.support.GenericApplicationContext;
 import org.springframework.test.context.MergedContextConfiguration;
+import org.springframework.test.context.aot.AotContextLoader;
 import org.springframework.util.StringUtils;
 
 /**
@@ -56,7 +58,7 @@ import org.springframework.util.StringUtils;
  * @since 2.5
  * @see #loadContext(MergedContextConfiguration)
  */
-public abstract class AbstractGenericContextLoader extends AbstractContextLoader {
+public abstract class AbstractGenericContextLoader extends AbstractContextLoader implements AotContextLoader {
 
 	protected static final Log logger = LogFactory.getLog(AbstractGenericContextLoader.class);
 
@@ -128,6 +130,27 @@ public abstract class AbstractGenericContextLoader extends AbstractContextLoader
 		return loadContext(mergedConfig, false);
 	}
 
+	@Override
+	public final ApplicationContext loadContextForAotRuntime(MergedContextConfiguration mergedConfig,
+			ApplicationContextInitializer<ConfigurableApplicationContext> initializer) throws Exception {
+
+		if (logger.isDebugEnabled()) {
+			logger.debug(String.format("Loading ApplicationContext for AOT runtime for merged context configuration [%s].",
+					mergedConfig));
+		}
+
+		validateMergedContextConfiguration(mergedConfig);
+
+		GenericApplicationContext context = createContext();
+		prepareContext(context);
+		prepareContext(context, mergedConfig);
+		initializer.initialize(context);
+		customizeContext(context);
+		customizeContext(context, mergedConfig);
+		context.refresh();
+		return context;
+	}
+
 	/**
 	 * Load a {@link GenericApplicationContext} for the supplied
 	 * {@link MergedContextConfiguration}.

spring-test/src/main/java/org/springframework/test/context/web/AbstractGenericWebContextLoader.java

@@ -22,6 +22,7 @@ import org.apache.commons.logging.LogFactory;
 
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextInitializer;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.annotation.AnnotationConfigUtils;
 import org.springframework.core.io.DefaultResourceLoader;
@@ -29,6 +30,7 @@ import org.springframework.core.io.FileSystemResourceLoader;
 import org.springframework.core.io.ResourceLoader;
 import org.springframework.mock.web.MockServletContext;
 import org.springframework.test.context.MergedContextConfiguration;
+import org.springframework.test.context.aot.AotContextLoader;
 import org.springframework.test.context.support.AbstractContextLoader;
 import org.springframework.util.Assert;
 import org.springframework.web.context.WebApplicationContext;
@@ -57,11 +59,60 @@ import org.springframework.web.context.support.GenericWebApplicationContext;
  * @see #loadContext(MergedContextConfiguration)
  * @see #loadContext(String...)
  */
-public abstract class AbstractGenericWebContextLoader extends AbstractContextLoader {
+public abstract class AbstractGenericWebContextLoader extends AbstractContextLoader implements AotContextLoader {
 
 	protected static final Log logger = LogFactory.getLog(AbstractGenericWebContextLoader.class);
 
 
+	// AotContextLoader
+
+	/**
+	 * Load a {@link GenericWebApplicationContext} for the supplied
+	 * {@link MergedContextConfiguration}.
+	 * <p>In contrast to {@link #loadContext(MergedContextConfiguration)}, this
+	 * method does not
+	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#refresh()
+	 * refresh} the {@code ApplicationContext} or
+	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#registerShutdownHook()
+	 * register a JVM shutdown hook} for it. Otherwise, this method implements
+	 * behavior identical to {@link #loadContext(MergedContextConfiguration)}.
+	 * @param mergedConfig the merged context configuration to use to load the
+	 * application context
+	 * @return a new web application context
+	 * @since 6.0
+	 * @see org.springframework.test.context.aot.AotContextLoader#loadContextForAotProcessing(MergedContextConfiguration)
+	 * @see GenericWebApplicationContext
+	 */
+	@Override
+	public final ApplicationContext loadContextForAotProcessing(
+			MergedContextConfiguration mergedConfig) throws Exception {
+
+		return loadContext(mergedConfig, false);
+	}
+
+	@Override
+	public final ApplicationContext loadContextForAotRuntime(MergedContextConfiguration mergedConfig,
+			ApplicationContextInitializer<ConfigurableApplicationContext> initializer) throws Exception {
+
+		if (!(mergedConfig instanceof WebMergedContextConfiguration webMergedConfig)) {
+			throw new IllegalArgumentException("""
+					Cannot load WebApplicationContext from non-web merged context configuration %s. \
+					Consider annotating your test class with @WebAppConfiguration."""
+						.formatted(mergedConfig));
+		}
+
+		validateMergedContextConfiguration(webMergedConfig);
+
+		GenericWebApplicationContext context = createContext();
+		configureWebResources(context, webMergedConfig);
+		prepareContext(context, webMergedConfig);
+		initializer.initialize(context);
+		customizeContext(context, webMergedConfig);
+		context.refresh();
+		return context;
+	}
+
+
 	// SmartContextLoader
 
 	/**
@@ -105,30 +156,6 @@ public abstract class AbstractGenericWebContextLoader extends AbstractContextLoa
 		return loadContext(mergedConfig, true);
 	}
 
-	/**
-	 * Load a {@link GenericWebApplicationContext} for the supplied
-	 * {@link MergedContextConfiguration}.
-	 * <p>In contrast to {@link #loadContext(MergedContextConfiguration)}, this
-	 * method does not
-	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#refresh()
-	 * refresh} the {@code ApplicationContext} or
-	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#registerShutdownHook()
-	 * register a JVM shutdown hook} for it. Otherwise, this method implements
-	 * behavior identical to {@link #loadContext(MergedContextConfiguration)}.
-	 * @param mergedConfig the merged context configuration to use to load the
-	 * application context
-	 * @return a new web application context
-	 * @since 6.0
-	 * @see org.springframework.test.context.SmartContextLoader#loadContextForAotProcessing(MergedContextConfiguration)
-	 * @see GenericWebApplicationContext
-	 */
-	@Override
-	public final ApplicationContext loadContextForAotProcessing(
-			MergedContextConfiguration mergedConfig) throws Exception {
-
-		return loadContext(mergedConfig, false);
-	}
-
 	/**
 	 * Load a {@link GenericWebApplicationContext} for the supplied
 	 * {@link MergedContextConfiguration}.

spring-test/src/test/java/org/springframework/test/context/support/DelegatingSmartContextLoaderTests.java

@@ -110,7 +110,7 @@ class DelegatingSmartContextLoaderTests {
 			assertThatIllegalStateException()
 			.isThrownBy(() -> loader.loadContext(mergedConfig))
 				.withMessageStartingWith("Neither")
-				.withMessageContaining("was able to load an ApplicationContext from");
+				.withMessageContaining("is able to load an ApplicationContext for");
 		}
 
 		/**