Update SpringBootContextLoader to support AOT

Update `SpringBootContextLoader` so that it now implements the
`AotContextLoader` interface. The `ContextLoaderHook` will abandon
at `contextLoaded` if the test class is being AOT processed.

This commit also introduces a new `AotApplicationContextInitializer`
which allows us to plug-in an alternative AOT application context
listener when the `SpringApplication` is running in test mode.

Closes gh-31965
This commit is contained in:
Phillip Webb
2022-09-12 13:29:48 -07:00
parent d1e7c9bd70
commit 4d037c3003
7 changed files with 469 additions and 11 deletions

View File

@@ -23,6 +23,7 @@ import java.util.Arrays;
import java.util.List;
import org.springframework.beans.BeanUtils;
import org.springframework.boot.AotApplicationContextInitializer;
import org.springframework.boot.ApplicationContextFactory;
import org.springframework.boot.ConfigurableBootstrapContext;
import org.springframework.boot.SpringApplication;
@@ -55,6 +56,8 @@ import org.springframework.test.context.ContextConfigurationAttributes;
import org.springframework.test.context.ContextCustomizer;
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.test.context.support.AbstractContextLoader;
import org.springframework.test.context.support.AnnotationConfigContextLoaderUtils;
import org.springframework.test.context.support.TestPropertySourceUtils;
@@ -90,15 +93,31 @@ import org.springframework.web.context.support.GenericWebApplicationContext;
* @since 1.4.0
* @see SpringBootTest
*/
public class SpringBootContextLoader extends AbstractContextLoader {
public class SpringBootContextLoader extends AbstractContextLoader implements AotContextLoader {
@Override
public ApplicationContext loadContext(MergedContextConfiguration mergedConfig) throws Exception {
return loadContext(mergedConfig, Mode.STANDARD, null);
}
@Override
public ApplicationContext loadContextForAotProcessing(MergedContextConfiguration mergedConfig) throws Exception {
return loadContext(mergedConfig, Mode.AOT_PROCESSING, null);
}
@Override
public ApplicationContext loadContextForAotRuntime(MergedContextConfiguration mergedConfig,
ApplicationContextInitializer<ConfigurableApplicationContext> initializer) throws Exception {
return loadContext(mergedConfig, Mode.AOT_RUNTIME, initializer);
}
private ApplicationContext loadContext(MergedContextConfiguration mergedConfig, Mode mode,
ApplicationContextInitializer<ConfigurableApplicationContext> initializer) {
assertHasClassesOrLocations(mergedConfig);
SpringBootTestAnnotation annotation = SpringBootTestAnnotation.get(mergedConfig);
String[] args = annotation.getArgs();
UseMainMethod useMainMethod = annotation.getUseMainMethod();
ContextLoaderHook hook = new ContextLoaderHook(mergedConfig);
ContextLoaderHook hook = new ContextLoaderHook(mergedConfig, mode, initializer);
if (useMainMethod != UseMainMethod.NEVER) {
Method mainMethod = getMainMethod(mergedConfig, useMainMethod);
if (mainMethod != null) {
@@ -297,6 +316,31 @@ public class SpringBootContextLoader extends AbstractContextLoader {
throw new IllegalStateException();
}
/**
* Modes that the {@link SpringBootContextLoader} can operate.
*/
private enum Mode {
/**
* Load for regular usage.
* @see SmartContextLoader#loadContext
*/
STANDARD,
/**
* Load for AOT processing.
* @see AotContextLoader#loadContextForAotProcessing
*/
AOT_PROCESSING,
/**
* Load for AOT runtime.
* @see AotContextLoader#loadContextForAotRuntime
*/
AOT_RUNTIME
}
/**
* Inner class to configure {@link WebMergedContextConfiguration}.
*/
@@ -417,8 +461,15 @@ public class SpringBootContextLoader extends AbstractContextLoader {
private final MergedContextConfiguration mergedConfig;
ContextLoaderHook(MergedContextConfiguration mergedConfig) {
private final Mode mode;
private final ApplicationContextInitializer<ConfigurableApplicationContext> initializer;
ContextLoaderHook(MergedContextConfiguration mergedConfig, Mode mode,
ApplicationContextInitializer<ConfigurableApplicationContext> initializer) {
this.mergedConfig = mergedConfig;
this.mode = mode;
this.initializer = initializer;
}
@Override
@@ -428,6 +479,17 @@ public class SpringBootContextLoader extends AbstractContextLoader {
@Override
public void starting(ConfigurableBootstrapContext bootstrapContext) {
SpringBootContextLoader.this.configure(ContextLoaderHook.this.mergedConfig, application);
if (ContextLoaderHook.this.initializer != null) {
application.addInitializers(
AotApplicationContextInitializer.of(ContextLoaderHook.this.initializer));
}
}
@Override
public void contextLoaded(ConfigurableApplicationContext context) {
if (ContextLoaderHook.this.mode == Mode.AOT_PROCESSING) {
throw new AbandonedRunException(context);
}
}
@Override

View File

@@ -49,6 +49,7 @@ import org.springframework.test.context.TestContext;
import org.springframework.test.context.TestContextAnnotationUtils;
import org.springframework.test.context.TestContextBootstrapper;
import org.springframework.test.context.TestExecutionListener;
import org.springframework.test.context.aot.AotTestAttributes;
import org.springframework.test.context.support.DefaultTestContextBootstrapper;
import org.springframework.test.context.support.TestPropertySourceUtils;
import org.springframework.test.context.web.WebAppConfiguration;
@@ -97,6 +98,16 @@ public class SpringBootTestContextBootstrapper extends DefaultTestContextBootstr
private static final Log logger = LogFactory.getLog(SpringBootTestContextBootstrapper.class);
private final AotTestAttributes aotTestAttributes;
public SpringBootTestContextBootstrapper() {
this(AotTestAttributes.getInstance());
}
SpringBootTestContextBootstrapper(AotTestAttributes aotTestAttributes) {
this.aotTestAttributes = aotTestAttributes;
}
@Override
public TestContext buildTestContext() {
TestContext context = super.buildTestContext();
@@ -231,14 +242,25 @@ public class SpringBootTestContextBootstrapper extends DefaultTestContextBootstr
if (containsNonTestComponent(classes) || mergedConfig.hasLocations()) {
return classes;
}
Class<?> found = new AnnotatedClassFinder(SpringBootConfiguration.class)
.findFromClass(mergedConfig.getTestClass());
Class<?> found = findConfigurationClass(mergedConfig.getTestClass());
Assert.state(found != null, "Unable to find a @SpringBootConfiguration, you need to use "
+ "@ContextConfiguration or @SpringBootTest(classes=...) with your test");
logger.info("Found @SpringBootConfiguration " + found.getName() + " for test " + mergedConfig.getTestClass());
return merge(found, classes);
}
private Class<?> findConfigurationClass(Class<?> testClass) {
String propertyName = "%s.SpringBootConfiguration.%s"
.formatted(SpringBootTestContextBootstrapper.class.getName(), testClass.getName());
String foundClassName = this.aotTestAttributes.getString(propertyName);
if (foundClassName != null) {
return ClassUtils.resolveClassName(foundClassName, testClass.getClassLoader());
}
Class<?> found = new AnnotatedClassFinder(SpringBootConfiguration.class).findFromClass(testClass);
this.aotTestAttributes.setAttribute(propertyName, found.getName());
return found;
}
private boolean containsNonTestComponent(Class<?>[] classes) {
for (Class<?> candidate : classes) {
if (!MergedAnnotations.from(candidate, SearchStrategy.INHERITED_ANNOTATIONS)