Spring/spring-boot
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update SpringBootContextLoader to support AOT

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

1
spring-boot-project/spring-boot-test/build.gradle
View File

@@ -55,6 +55,7 @@ dependencies {
testImplementation("org.slf4j:slf4j-api")
testImplementation("org.spockframework:spock-core")
testImplementation("org.springframework:spring-webmvc")
testImplementation("org.springframework:spring-core-test")
testImplementation("org.testng:testng")
testRuntimeOnly("org.junit.vintage:junit-vintage-engine")

68
spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/SpringBootContextLoader.java
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

26
spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/SpringBootTestContextBootstrapper.java
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)

112
spring-boot-project/spring-boot-test/src/test/java/org/springframework/boot/test/context/SpringBootContextLoaderAotTests.java Normal file
View File

@@ -0,0 +1,112 @@
/*
* Copyright 2012-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.boot.test.context;
import java.util.stream.Stream;
import org.junit.jupiter.api.Test;
import org.springframework.aot.AotDetector;
import org.springframework.aot.generate.InMemoryGeneratedFiles;
import org.springframework.aot.test.generate.compile.CompileWithTargetClassAccess;
import org.springframework.aot.test.generate.compile.TestCompiler;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest.WebEnvironment;
import org.springframework.context.ApplicationContextInitializer;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.annotation.Import;
import org.springframework.context.support.GenericApplicationContext;
import org.springframework.test.context.BootstrapUtils;
import org.springframework.test.context.MergedContextConfiguration;
import org.springframework.test.context.TestContextBootstrapper;
import org.springframework.test.context.aot.AotContextLoader;
import org.springframework.test.context.aot.AotTestContextInitializers;
import org.springframework.test.context.aot.TestContextAotGenerator;
import org.springframework.test.util.ReflectionTestUtils;
import org.springframework.util.ClassUtils;
import org.springframework.util.function.ThrowingConsumer;
import static org.assertj.core.api.Assertions.assertThat;
/**
* Tests for {@link SpringBootContextLoader} when used in AOT mode.
*
* @author Phillip Webb
*/
@CompileWithTargetClassAccess
class SpringBootContextLoaderAotTests {
@Test
void loadContextForAotProcessingAndAotRuntime() {
InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
TestContextAotGenerator generator = new TestContextAotGenerator(generatedFiles);
Class<?> testClass = ExampleTest.class;
generator.processAheadOfTime(Stream.of(testClass));
TestCompiler.forSystem().withFiles(generatedFiles).printFiles(System.out)
.compile(ThrowingConsumer.of((compiled) -> assertCompiledTest(testClass)));
}
private void assertCompiledTest(Class<?> testClass) throws Exception {
try {
System.setProperty(AotDetector.AOT_ENABLED, "true");
resetAotClasses();
AotTestContextInitializers aotContextInitializers = new AotTestContextInitializers();
TestContextBootstrapper testContextBootstrapper = BootstrapUtils.resolveTestContextBootstrapper(testClass);
MergedContextConfiguration mergedConfig = testContextBootstrapper.buildMergedContextConfiguration();
ApplicationContextInitializer<ConfigurableApplicationContext> contextInitializer = aotContextInitializers
.getContextInitializer(testClass);
ConfigurableApplicationContext context = (ConfigurableApplicationContext) ((AotContextLoader) mergedConfig
.getContextLoader()).loadContextForAotRuntime(mergedConfig, contextInitializer);
assertThat(context).isExactlyInstanceOf(GenericApplicationContext.class);
String[] beanNames = context.getBeanNamesForType(ExampleBean.class);
BeanDefinition beanDefinition = context.getBeanFactory().getBeanDefinition(beanNames[0]);
assertThat(beanDefinition).isNotExactlyInstanceOf(GenericBeanDefinition.class);
}
finally {
System.clearProperty(AotDetector.AOT_ENABLED);
resetAotClasses();
}
}
private void resetAotClasses() {
reset("org.springframework.test.context.aot.AotTestAttributesFactory");
reset("org.springframework.test.context.aot.AotTestContextInitializersFactory");
}
private void reset(String className) {
Class<?> targetClass = ClassUtils.resolveClassName(className, null);
ReflectionTestUtils.invokeMethod(targetClass, "reset");
}
@SpringBootTest(classes = ExampleConfig.class, webEnvironment = WebEnvironment.NONE)
static class ExampleTest {
}
@SpringBootConfiguration
@Import(ExampleBean.class)
static class ExampleConfig {
}
static class ExampleBean {
}
}

152
spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotApplicationContextInitializer.java Normal file
View File

@@ -0,0 +1,152 @@
/*
* Copyright 2012-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.boot;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.BeanUtils;
import org.springframework.context.ApplicationContextInitializer;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.core.log.LogMessage;
import org.springframework.util.Assert;
import org.springframework.util.ClassUtils;
/**
* A {@link ApplicationContextInitializer} wrapper used to initialize a
* {@link ConfigurableApplicationContext} using artifacts that were generated
* ahead-of-time.
*
* @param <C> the application context type
* @author Phillip Webb
* @since 3.0.0
*/
public abstract sealed class AotApplicationContextInitializer<C extends ConfigurableApplicationContext>
implements ApplicationContextInitializer<C> {
private static final Log logger = LogFactory.getLog(AotApplicationContextInitializer.class);
@Override
public final void initialize(C applicationContext) {
logger.debug(LogMessage.format("Initializing ApplicationContext using AOT initializer '%s'", getName()));
aotInitialize(applicationContext);
}
abstract void aotInitialize(C applicationContext);
abstract String getName();
static <C extends ConfigurableApplicationContext> AotApplicationContextInitializer<C> forMainApplicationClass(
Class<?> mainApplicationClass) {
String initializerClassName = mainApplicationClass.getName() + "__ApplicationContextInitializer";
return new ReflectionDelegatingAotApplicationContextInitializer<>(initializerClassName);
}
/**
* Create a new {@link AotApplicationContextInitializer} by delegating to an existing
* initializer instance.
* @param <C> the application context type
* @param initializer the initializer to delegate to
* @return a new {@link AotApplicationContextInitializer} instance
*/
public static <C extends ConfigurableApplicationContext> AotApplicationContextInitializer<C> of(
ApplicationContextInitializer<C> initializer) {
Assert.notNull(initializer, "Initializer must not be null");
return new InstanceDelegatingAotApplicationContextInitializer<>(initializer.getClass().getName(), initializer);
}
/**
* Create a new {@link AotApplicationContextInitializer} by delegating to an existing
* initializer instance.
* @param <C> the application context type
* @param initializer the initializer to delegate to
* @param name the name of the initializer
* @return a new {@link AotApplicationContextInitializer} instance
*/
public static <C extends ConfigurableApplicationContext> AotApplicationContextInitializer<C> of(String name,
ApplicationContextInitializer<C> initializer) {
Assert.notNull(name, "Name must not be null");
Assert.notNull(initializer, "Initializer must not be null");
return new InstanceDelegatingAotApplicationContextInitializer<>(name, initializer);
}
/**
* {@link AotApplicationContextInitializer} that delegates to an initializer created
* via reflection.
*
* @param <C> the application context type
*/
static final class ReflectionDelegatingAotApplicationContextInitializer<C extends ConfigurableApplicationContext>
extends AotApplicationContextInitializer<C> {
private final String initializerClassName;
ReflectionDelegatingAotApplicationContextInitializer(String initializerClassName) {
this.initializerClassName = initializerClassName;
}
@Override
void aotInitialize(C applicationContext) {
ApplicationContextInitializer<C> initializer = createInitializer(applicationContext.getClassLoader());
initializer.initialize(applicationContext);
}
@SuppressWarnings("unchecked")
private ApplicationContextInitializer<C> createInitializer(ClassLoader classLoader) {
Class<?> initializerClass = ClassUtils.resolveClassName(this.initializerClassName, classLoader);
Assert.isAssignable(ApplicationContextInitializer.class, initializerClass);
return (ApplicationContextInitializer<C>) BeanUtils.instantiateClass(initializerClass);
}
@Override
String getName() {
return this.initializerClassName;
}
}
/**
* {@link AotApplicationContextInitializer} that delegates to an existing initializer
* instance.
*
* @param <C> the application context type
*/
static final class InstanceDelegatingAotApplicationContextInitializer<C extends ConfigurableApplicationContext>
extends AotApplicationContextInitializer<C> {
private final String name;
private final ApplicationContextInitializer<C> initializer;
InstanceDelegatingAotApplicationContextInitializer(String name, ApplicationContextInitializer<C> initializer) {
this.name = name;
this.initializer = initializer;
}
@Override
void aotInitialize(C applicationContext) {
this.initializer.initialize(applicationContext);
}
@Override
String getName() {
return this.name;
}
}
}

22
spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java
View File

@@ -67,7 +67,6 @@ import org.springframework.context.annotation.AnnotationConfigApplicationContext
import org.springframework.context.annotation.AnnotationConfigUtils;
import org.springframework.context.annotation.ClassPathBeanDefinitionScanner;
import org.springframework.context.annotation.ConfigurationClassPostProcessor;
import org.springframework.context.aot.ApplicationContextAotInitializer;
import org.springframework.context.support.AbstractApplicationContext;
import org.springframework.context.support.GenericApplicationContext;
import org.springframework.core.GenericTypeResolver;
@@ -86,6 +85,7 @@ import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.ResourceLoader;
import org.springframework.core.io.support.SpringFactoriesLoader;
import org.springframework.core.io.support.SpringFactoriesLoader.ArgumentResolver;
import org.springframework.core.log.LogMessage;
import org.springframework.core.metrics.ApplicationStartup;
import org.springframework.util.Assert;
import org.springframework.util.ClassUtils;
@@ -416,12 +416,22 @@ public class SpringApplication {
private void addAotGeneratedInitializerIfNecessary(List<ApplicationContextInitializer<?>> initializers) {
if (AotDetector.useGeneratedArtifacts()) {
String initializerClassName = this.mainApplicationClass.getName() + "__ApplicationContextInitializer";
if (logger.isDebugEnabled()) {
logger.debug("Using AOT generated initializer: " + initializerClassName);
List<AotApplicationContextInitializer<?>> aotInitializers = new ArrayList<>();
for (ApplicationContextInitializer<?> candidate : initializers) {
if (candidate instanceof AotApplicationContextInitializer<?> aotInitializer) {
aotInitializers.add(aotInitializer);
}
}
initializers.add(0,
(context) -> new ApplicationContextAotInitializer().initialize(context, initializerClassName));
if (aotInitializers.isEmpty()) {
AotApplicationContextInitializer<?> aotInitializer = AotApplicationContextInitializer
.forMainApplicationClass(this.mainApplicationClass);
aotInitializers.add(aotInitializer);
}
for (AotApplicationContextInitializer<?> aotInitializer : aotInitializers) {
logger.debug(LogMessage.format("Using AOT generated initializer: %s", aotInitializer.getName()));
}
initializers.removeAll(aotInitializers);
initializers.addAll(0, aotInitializers);
}
}

99
spring-boot-project/spring-boot/src/test/java/org/springframework/boot/AotApplicationContextInitializerTests.java Normal file
View File

@@ -0,0 +1,99 @@
/*
* Copyright 2012-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.boot;
import org.junit.jupiter.api.Test;
import org.springframework.context.ApplicationContextInitializer;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.support.GenericApplicationContext;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
/**
* Tests for {@link AotApplicationContextInitializer}.
*
* @author Phillip Webb
*/
class AotApplicationContextInitializerTests {
@Test
void forMainApplicationClassUsesReflection() {
AotApplicationContextInitializer<ConfigurableApplicationContext> initializer = AotApplicationContextInitializer
.forMainApplicationClass(ExampleMain.class);
GenericApplicationContext applicationContext = new GenericApplicationContext();
initializer.initialize(applicationContext);
assertThat(initializer.getName()).isEqualTo(ExampleMain__ApplicationContextInitializer.class.getName());
assertThat(applicationContext.getId()).isEqualTo("ExampleMain");
}
@Test
void ofWhenInitializerIsNullThrowsException() {
assertThatIllegalArgumentException().isThrownBy(() -> AotApplicationContextInitializer.of(null))
.withMessage("Initializer must not be null");
}
@Test
void ofWithNameWhenInitializerIsNullThrowsException() {
assertThatIllegalArgumentException().isThrownBy(() -> AotApplicationContextInitializer.of("test", null))
.withMessage("Initializer must not be null");
}
@Test
void ofWithNameWhenNameIsNullThrowsException() {
assertThatIllegalArgumentException()
.isThrownBy(() -> AotApplicationContextInitializer.of(null, new CustomApplicationContextInitializer()))
.withMessage("Name must not be null");
}
@Test
void ofUsesDelegation() {
AotApplicationContextInitializer<ConfigurableApplicationContext> initializer = AotApplicationContextInitializer
.of(new CustomApplicationContextInitializer());
GenericApplicationContext applicationContext = new GenericApplicationContext();
initializer.initialize(applicationContext);
assertThat(initializer.getName()).isEqualTo(CustomApplicationContextInitializer.class.getName());
assertThat(applicationContext.getId()).isEqualTo("Custom");
}
static class ExampleMain {
}
static class ExampleMain__ApplicationContextInitializer
implements ApplicationContextInitializer<ConfigurableApplicationContext> {
@Override
public void initialize(ConfigurableApplicationContext applicationContext) {
applicationContext.setId("ExampleMain");
}
}
static class CustomApplicationContextInitializer
implements ApplicationContextInitializer<ConfigurableApplicationContext> {
@Override
public void initialize(ConfigurableApplicationContext applicationContext) {
applicationContext.setId("Custom");
}
}
}