Skip to content

Commit ada0880

Browse files
sbrannensbrannen
committed
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
1 parent b0d6570 commit ada0880

File tree

6 files changed

+315
-46
lines changed

6 files changed

+315
-46
lines changed

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

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
/*
2+
* Copyright 2002-2022 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.test.context.aot;
18+
19+
import java.lang.reflect.Method;
20+
import java.util.Map;
21+
import java.util.function.Supplier;
22+
23+
import org.springframework.context.ApplicationContextInitializer;
24+
import org.springframework.context.ConfigurableApplicationContext;
25+
import org.springframework.util.Assert;
26+
import org.springframework.util.ClassUtils;
27+
import org.springframework.util.ReflectionUtils;
28+
29+
/**
30+
* {@code AotTestMappings} provides mappings from test classes to AOT-optimized
31+
* context initializers.
32+
*
33+
* <p>If a test class is not {@linkplain #isSupportedTestClass(Class) supported} in
34+
* AOT mode, {@link #getContextInitializer(Class)} will return {@code null}.
35+
*
36+
* <p>Reflectively accesses {@link #GENERATED_MAPPINGS_CLASS_NAME} generated by
37+
* the {@link TestContextAotGenerator} to retrieve the mappings generated during
38+
* AOT processing.
39+
*
40+
* @author Sam Brannen
41+
* @author Stephane Nicoll
42+
* @since 6.0
43+
*/
44+
public class AotTestMappings {
45+
46+
// TODO Add support in ClassNameGenerator for supplying a predefined class name.
47+
// There is a similar issue in Spring Boot where code relies on a generated name.
48+
// Ideally we would generate a class named: org.springframework.test.context.aot.GeneratedAotTestMappings
49+
static final String GENERATED_MAPPINGS_CLASS_NAME = AotTestMappings.class.getName() + "__Generated";
50+
51+
static final String GENERATED_MAPPINGS_METHOD_NAME = "getContextInitializers";
52+
53+
private final Map<String, Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>>> contextInitializers;
54+
55+
56+
public AotTestMappings() {
57+
this(GENERATED_MAPPINGS_CLASS_NAME);
58+
}
59+
60+
AotTestMappings(String initializerClassName) {
61+
this(loadContextInitializersMap(initializerClassName));
62+
}
63+
64+
AotTestMappings(Map<String, Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>>> contextInitializers) {
65+
this.contextInitializers = contextInitializers;
66+
}
67+
68+
69+
/**
70+
* Determine if the specified test class has an AOT-optimized application context
71+
* initializer.
72+
* <p>If this method returns {@code true}, {@link #getContextInitializer(Class)}
73+
* should not return {@code null}.
74+
*/
75+
public boolean isSupportedTestClass(Class<?> testClass) {
76+
return this.contextInitializers.containsKey(testClass.getName());
77+
}
78+
79+
/**
80+
* Get the AOT {@link ApplicationContextInitializer} for the specified test class.
81+
* @return the AOT context initializer, or {@code null} if there is no AOT context
82+
* initializer for the specified test class
83+
* @see #isSupportedTestClass(Class)
84+
*/
85+
public ApplicationContextInitializer<ConfigurableApplicationContext> getContextInitializer(Class<?> testClass) {
86+
Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>> supplier =
87+
this.contextInitializers.get(testClass.getName());
88+
return (supplier != null ? supplier.get() : null);
89+
}
90+
91+
92+
@SuppressWarnings({ "rawtypes", "unchecked" })
93+
private static Map<String, Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>>>
94+
loadContextInitializersMap(String className) {
95+
96+
String methodName = GENERATED_MAPPINGS_METHOD_NAME;
97+
98+
try {
99+
Class<?> clazz = ClassUtils.forName(className, null);
100+
Method method = ReflectionUtils.findMethod(clazz, methodName);
101+
Assert.state(method != null, () -> "No %s() method found in %s".formatted(methodName, clazz.getName()));
102+
return (Map<String, Supplier<ApplicationContextInitializer<ConfigurableApplicationContext>>>)
103+
ReflectionUtils.invokeMethod(method, null);
104+
}
105+
catch (IllegalStateException ex) {
106+
throw ex;
107+
}
108+
catch (Exception ex) {
109+
throw new IllegalStateException("Failed to invoke %s() method in %s".formatted(methodName, className), ex);
110+
}
111+
}
112+
113+
}

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

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
/*
2+
* Copyright 2002-2022 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.test.context.aot;
18+
19+
import java.util.HashMap;
20+
import java.util.List;
21+
import java.util.Map;
22+
import java.util.function.Supplier;
23+
24+
import javax.lang.model.element.Modifier;
25+
26+
import org.apache.commons.logging.Log;
27+
import org.apache.commons.logging.LogFactory;
28+
29+
import org.springframework.aot.generate.GeneratedClass;
30+
import org.springframework.aot.generate.GeneratedClasses;
31+
import org.springframework.context.ApplicationContextInitializer;
32+
import org.springframework.context.ConfigurableApplicationContext;
33+
import org.springframework.core.log.LogMessage;
34+
import org.springframework.javapoet.ClassName;
35+
import org.springframework.javapoet.CodeBlock;
36+
import org.springframework.javapoet.MethodSpec;
37+
import org.springframework.javapoet.ParameterizedTypeName;
38+
import org.springframework.javapoet.TypeName;
39+
import org.springframework.javapoet.TypeSpec;
40+
import org.springframework.javapoet.WildcardTypeName;
41+
import org.springframework.util.MultiValueMap;
42+
43+
/**
44+
* Internal code generator for mappings used by {@link AotTestMappings}.
45+
*
46+
* @author Sam Brannen
47+
* @since 6.0
48+
*/
49+
class AotTestMappingsCodeGenerator {
50+
51+
private static final Log logger = LogFactory.getLog(AotTestMappingsCodeGenerator.class);
52+
53+
private static final ParameterizedTypeName CONTEXT_INITIALIZER = ParameterizedTypeName.get(
54+
ClassName.get(ApplicationContextInitializer.class),
55+
WildcardTypeName.subtypeOf(ConfigurableApplicationContext.class));
56+
57+
private static final ParameterizedTypeName CONTEXT_INITIALIZER_SUPPLIER = ParameterizedTypeName
58+
.get(ClassName.get(Supplier.class), CONTEXT_INITIALIZER);
59+
60+
// Map<String, Supplier<ApplicationContextInitializer<? extends ConfigurableApplicationContext>>>
61+
private static final TypeName CONTEXT_SUPPLIER_MAP = ParameterizedTypeName
62+
.get(ClassName.get(Map.class), ClassName.get(String.class), CONTEXT_INITIALIZER_SUPPLIER);
63+
64+
65+
private final MultiValueMap<ClassName, Class<?>> initializerClassMappings;
66+
67+
private final GeneratedClass generatedClass;
68+
69+
70+
AotTestMappingsCodeGenerator(MultiValueMap<ClassName, Class<?>> initializerClassMappings,
71+
GeneratedClasses generatedClasses) {
72+
73+
this.initializerClassMappings = initializerClassMappings;
74+
this.generatedClass = generatedClasses.addForFeature("Generated", this::generateType);
75+
}
76+
77+
78+
GeneratedClass getGeneratedClass() {
79+
return this.generatedClass;
80+
}
81+
82+
private void generateType(TypeSpec.Builder type) {
83+
logger.debug(LogMessage.format("Generating AOT test mappings in %s",
84+
this.generatedClass.getName().reflectionName()));
85+
type.addJavadoc("Generated mappings for {@link $T}.", AotTestMappings.class);
86+
type.addModifiers(Modifier.PUBLIC);
87+
type.addMethod(generateMappingMethod());
88+
}
89+
90+
private MethodSpec generateMappingMethod() {
91+
MethodSpec.Builder method = MethodSpec.methodBuilder(AotTestMappings.GENERATED_MAPPINGS_METHOD_NAME);
92+
method.addModifiers(Modifier.PUBLIC, Modifier.STATIC);
93+
method.returns(CONTEXT_SUPPLIER_MAP);
94+
method.addCode(generateMappingCode());
95+
return method.build();
96+
}
97+
98+
private CodeBlock generateMappingCode() {
99+
CodeBlock.Builder code = CodeBlock.builder();
100+
code.addStatement("$T map = new $T<>()", CONTEXT_SUPPLIER_MAP, HashMap.class);
101+
this.initializerClassMappings.forEach((className, testClasses) -> {
102+
List<String> testClassNames = testClasses.stream().map(Class::getName).toList();
103+
logger.debug(LogMessage.format(
104+
"Generating mapping from AOT context initializer [%s] to test classes %s",
105+
className.reflectionName(), testClassNames));
106+
testClassNames.forEach(testClassName ->
107+
code.addStatement("map.put($S, () -> new $T())", testClassName, className));
108+
});
109+
code.addStatement("return map");
110+
return code.build();
111+
}
112+
113+
}

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

Lines changed: 43 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,17 @@
2424

2525
import org.springframework.aot.generate.ClassNameGenerator;
2626
import org.springframework.aot.generate.DefaultGenerationContext;
27+
import org.springframework.aot.generate.GeneratedClasses;
2728
import org.springframework.aot.generate.GeneratedFiles;
2829
import org.springframework.aot.generate.GenerationContext;
30+
import org.springframework.aot.hint.MemberCategory;
2931
import org.springframework.aot.hint.RuntimeHints;
32+
import org.springframework.aot.hint.TypeReference;
3033
import org.springframework.context.ApplicationContext;
3134
import org.springframework.context.ApplicationContextInitializer;
3235
import org.springframework.context.aot.ApplicationContextAotGenerator;
3336
import org.springframework.context.support.GenericApplicationContext;
37+
import org.springframework.core.log.LogMessage;
3438
import org.springframework.javapoet.ClassName;
3539
import org.springframework.test.context.BootstrapUtils;
3640
import org.springframework.test.context.ContextLoader;
@@ -51,7 +55,7 @@
5155
*/
5256
class TestContextAotGenerator {
5357

54-
private static final Log logger = LogFactory.getLog(TestClassScanner.class);
58+
private static final Log logger = LogFactory.getLog(TestContextAotGenerator.class);
5559

5660
private final ApplicationContextAotGenerator aotGenerator = new ApplicationContextAotGenerator();
5761

@@ -97,30 +101,33 @@ public final RuntimeHints getRuntimeHints() {
97101
* @throws TestContextAotException if an error occurs during AOT processing
98102
*/
99103
public void processAheadOfTime(Stream<Class<?>> testClasses) throws TestContextAotException {
100-
MultiValueMap<MergedContextConfiguration, Class<?>> map = new LinkedMultiValueMap<>();
101-
testClasses.forEach(testClass -> map.add(buildMergedContextConfiguration(testClass), testClass));
102-
103-
map.forEach((mergedConfig, classes) -> {
104-
// System.err.println(mergedConfig + " -> " + classes);
105-
if (logger.isDebugEnabled()) {
106-
logger.debug("Generating AOT artifacts for test classes [%s]"
107-
.formatted(classes.stream().map(Class::getCanonicalName).toList()));
108-
}
104+
MultiValueMap<MergedContextConfiguration, Class<?>> mergedConfigMappings = new LinkedMultiValueMap<>();
105+
testClasses.forEach(testClass -> mergedConfigMappings.add(buildMergedContextConfiguration(testClass), testClass));
106+
processAheadOfTime(mergedConfigMappings);
107+
}
108+
109+
private void processAheadOfTime(MultiValueMap<MergedContextConfiguration, Class<?>> mergedConfigMappings) {
110+
MultiValueMap<ClassName, Class<?>> initializerClassMappings = new LinkedMultiValueMap<>();
111+
mergedConfigMappings.forEach((mergedConfig, testClasses) -> {
112+
logger.debug(LogMessage.format("Generating AOT artifacts for test classes %s",
113+
testClasses.stream().map(Class::getName).toList()));
109114
try {
110115
// Use first test class discovered for a given unique MergedContextConfiguration.
111-
Class<?> testClass = classes.get(0);
116+
Class<?> testClass = testClasses.get(0);
112117
DefaultGenerationContext generationContext = createGenerationContext(testClass);
113-
ClassName className = processAheadOfTime(mergedConfig, generationContext);
114-
// TODO Store ClassName in a map analogous to TestContextAotProcessor in Spring Native.
118+
ClassName initializer = processAheadOfTime(mergedConfig, generationContext);
119+
Assert.state(!initializerClassMappings.containsKey(initializer),
120+
() -> "ClassName [%s] already encountered".formatted(initializer.reflectionName()));
121+
initializerClassMappings.addAll(initializer, testClasses);
115122
generationContext.writeGeneratedContent();
116123
}
117124
catch (Exception ex) {
118-
if (logger.isWarnEnabled()) {
119-
logger.warn("Failed to generate AOT artifacts for test classes [%s]"
120-
.formatted(classes.stream().map(Class::getCanonicalName).toList()), ex);
121-
}
125+
logger.warn(LogMessage.format("Failed to generate AOT artifacts for test classes [%s]",
126+
testClasses.stream().map(Class::getName).toList()), ex);
122127
}
123128
});
129+
130+
generateAotTestMappings(initializerClassMappings);
124131
}
125132

126133
/**
@@ -143,7 +150,7 @@ ClassName processAheadOfTime(MergedContextConfiguration mergedConfig,
143150
}
144151
catch (Throwable ex) {
145152
throw new TestContextAotException("Failed to process test class [%s] for AOT"
146-
.formatted(mergedConfig.getTestClass().getCanonicalName()), ex);
153+
.formatted(mergedConfig.getTestClass().getName()), ex);
147154
}
148155
}
149156

@@ -154,7 +161,7 @@ ClassName processAheadOfTime(MergedContextConfiguration mergedConfig,
154161
* create {@link GenericApplicationContext GenericApplicationContexts}.
155162
* @throws TestContextAotException if an error occurs while loading the application
156163
* context or if one of the prerequisites is not met
157-
* @see SmartContextLoader#loadContextForAotProcessing(MergedContextConfiguration)
164+
* @see AotContextLoader#loadContextForAotProcessing(MergedContextConfiguration)
158165
*/
159166
private GenericApplicationContext loadContextForAotProcessing(
160167
MergedContextConfiguration mergedConfig) throws TestContextAotException {
@@ -164,7 +171,7 @@ private GenericApplicationContext loadContextForAotProcessing(
164171
Assert.notNull(contextLoader, """
165172
Cannot load an ApplicationContext with a NULL 'contextLoader'. \
166173
Consider annotating test class [%s] with @ContextConfiguration or \
167-
@ContextHierarchy.""".formatted(testClass.getCanonicalName()));
174+
@ContextHierarchy.""".formatted(testClass.getName()));
168175

169176
if (contextLoader instanceof AotContextLoader aotContextLoader) {
170177
try {
@@ -176,13 +183,13 @@ private GenericApplicationContext loadContextForAotProcessing(
176183
catch (Exception ex) {
177184
throw new TestContextAotException(
178185
"Failed to load ApplicationContext for AOT processing for test class [%s]"
179-
.formatted(testClass.getCanonicalName()), ex);
186+
.formatted(testClass.getName()), ex);
180187
}
181188
}
182189
throw new TestContextAotException("""
183190
Cannot generate AOT artifacts for test class [%s]. The configured \
184191
ContextLoader [%s] must be an AotContextLoader and must create a \
185-
GenericApplicationContext.""".formatted(testClass.getCanonicalName(),
192+
GenericApplicationContext.""".formatted(testClass.getName(),
186193
contextLoader.getClass().getName()));
187194
}
188195

@@ -203,4 +210,18 @@ private String nextTestContextId() {
203210
return "TestContext%03d_".formatted(this.sequence.incrementAndGet());
204211
}
205212

213+
private void generateAotTestMappings(MultiValueMap<ClassName, Class<?>> initializerClassMappings) {
214+
ClassNameGenerator classNameGenerator = new ClassNameGenerator(AotTestMappings.class);
215+
DefaultGenerationContext generationContext =
216+
new DefaultGenerationContext(classNameGenerator, this.generatedFiles, this.runtimeHints);
217+
GeneratedClasses generatedClasses = generationContext.getGeneratedClasses();
218+
219+
AotTestMappingsCodeGenerator codeGenerator =
220+
new AotTestMappingsCodeGenerator(initializerClassMappings, generatedClasses);
221+
generationContext.writeGeneratedContent();
222+
String className = codeGenerator.getGeneratedClass().getName().reflectionName();
223+
this.runtimeHints.reflection().registerType(TypeReference.of(className),
224+
builder -> builder.withMembers(MemberCategory.INVOKE_PUBLIC_METHODS));
225+
}
226+
206227
}

‎spring-test/src/test/java/org/springframework/test/context/aot/AbstractAotTests.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@
2828
abstract class AbstractAotTests {
2929

3030
static final String[] expectedSourceFilesForBasicSpringTests = {
31+
// Global
32+
"org/springframework/test/context/aot/AotTestMappings__Generated.java",
3133
// BasicSpringJupiterSharedConfigTests
3234
"org/springframework/context/event/DefaultEventListenerFactory__TestContext001_BeanDefinitions.java",
3335
"org/springframework/context/event/EventListenerMethodProcessor__TestContext001_BeanDefinitions.java",

‎spring-test/src/test/java/org/springframework/test/context/aot/AotSmokeTests.java

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,9 +49,11 @@ void scanClassPathThenGenerateSourceFilesAndCompileThem() {
4949
List<String> sourceFiles = generatedFiles.getGeneratedFiles(Kind.SOURCE).keySet().stream().toList();
5050
assertThat(sourceFiles).containsExactlyInAnyOrder(expectedSourceFilesForBasicSpringTests);
5151

52-
TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled -> {
53-
// just make sure compilation completes without errors
54-
});
52+
TestCompiler.forSystem().withFiles(generatedFiles)
53+
// .printFiles(System.out)
54+
.compile(compiled -> {
55+
// just make sure compilation completes without errors
56+
});
5557
}
5658

5759
}

0 commit comments

Comments
 (0)