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
This commit is contained in:
Sam Brannen
2022-08-19 14:21:57 +02:00
parent b0d65709ac
commit ada0880f3c
6 changed files with 314 additions and 45 deletions

View File

@@ -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);
}
}
}

View File

@@ -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();
}
}

View File

@@ -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));
MultiValueMap<MergedContextConfiguration, Class<?>> mergedConfigMappings = new LinkedMultiValueMap<>();
testClasses.forEach(testClass -> mergedConfigMappings.add(buildMergedContextConfiguration(testClass), testClass));
processAheadOfTime(mergedConfigMappings);
}
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()));
}
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));
}
}