Skip to content

Commit

Permalink
Introduce AotTestMappings and AotTestMappingsCodeGenerator
Browse files Browse the repository at this point in the history
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 spring-projectsgh-28205
Closes spring-projectsgh-28204
  • Loading branch information
sbrannen committed Aug 23, 2022
1 parent b0d6570 commit ada0880
Show file tree
Hide file tree
Showing 6 changed files with 315 additions and 46 deletions.
Original file line number Diff line number Diff line change
@@ -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);
}
}

}
Original file line number Diff line number Diff line change
@@ -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();
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,17 @@

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;
Expand All @@ -51,7 +55,7 @@
*/
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();

Expand Down Expand Up @@ -97,30 +101,33 @@ public final RuntimeHints getRuntimeHints() {
* @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));

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

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

/**
Expand All @@ -143,7 +150,7 @@ ClassName processAheadOfTime(MergedContextConfiguration mergedConfig,
}
catch (Throwable ex) {
throw new TestContextAotException("Failed to process test class [%s] for AOT"
.formatted(mergedConfig.getTestClass().getCanonicalName()), ex);
.formatted(mergedConfig.getTestClass().getName()), ex);
}
}

Expand All @@ -154,7 +161,7 @@ ClassName processAheadOfTime(MergedContextConfiguration mergedConfig,
* 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 {
Expand All @@ -164,7 +171,7 @@ private GenericApplicationContext loadContextForAotProcessing(
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 {
Expand All @@ -176,13 +183,13 @@ private GenericApplicationContext loadContextForAotProcessing(
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()));
}

Expand All @@ -203,4 +210,18 @@ private String nextTestContextId() {
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));
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@
abstract class AbstractAotTests {

static final String[] expectedSourceFilesForBasicSpringTests = {
// Global
"org/springframework/test/context/aot/AotTestMappings__Generated.java",
// BasicSpringJupiterSharedConfigTests
"org/springframework/context/event/DefaultEventListenerFactory__TestContext001_BeanDefinitions.java",
"org/springframework/context/event/EventListenerMethodProcessor__TestContext001_BeanDefinitions.java",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,9 +49,11 @@ void scanClassPathThenGenerateSourceFilesAndCompileThem() {
List<String> sourceFiles = generatedFiles.getGeneratedFiles(Kind.SOURCE).keySet().stream().toList();
assertThat(sourceFiles).containsExactlyInAnyOrder(expectedSourceFilesForBasicSpringTests);

TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled -> {
// just make sure compilation completes without errors
});
TestCompiler.forSystem().withFiles(generatedFiles)
// .printFiles(System.out)
.compile(compiled -> {
// just make sure compilation completes without errors
});
}

}
Loading

0 comments on commit ada0880

Please sign in to comment.