Skip to content

Commit 88a60ec

Browse files
nosannosan
committed
Introduce ImportTestContainersBeanFactoryInitializationAotProcessor that collects all importing classes and then generates an initializer method that invokes ImportTestcontainersRegistrar.registerBeanDefinitions(...) for those classes
1 parent feb8abf commit 88a60ec

File tree

4 files changed

+143
-262
lines changed

4 files changed

+143
-262
lines changed

‎spring-boot-project/spring-boot-testcontainers/src/main/java/org/springframework/boot/testcontainers/context/DynamicPropertySourceMethodsImporter.java

Lines changed: 0 additions & 162 deletions
Original file line numberDiff line numberDiff line change
@@ -18,37 +18,16 @@
1818

1919
import java.lang.reflect.Method;
2020
import java.lang.reflect.Modifier;
21-
import java.util.Map;
2221
import java.util.Set;
2322

24-
import org.springframework.aot.generate.AccessControl;
25-
import org.springframework.aot.generate.GeneratedClass;
26-
import org.springframework.aot.generate.GeneratedMethod;
27-
import org.springframework.aot.generate.GenerationContext;
28-
import org.springframework.aot.generate.MethodReference.ArgumentCodeGenerator;
29-
import org.springframework.aot.hint.ExecutableMode;
30-
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution;
31-
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor;
32-
import org.springframework.beans.factory.aot.BeanFactoryInitializationCode;
33-
import org.springframework.beans.factory.aot.BeanRegistrationExcludeFilter;
34-
import org.springframework.beans.factory.config.BeanDefinition;
35-
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
3623
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
37-
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
38-
import org.springframework.beans.factory.support.RegisteredBean;
39-
import org.springframework.beans.factory.support.RootBeanDefinition;
4024
import org.springframework.boot.testcontainers.properties.TestcontainersPropertySource;
4125
import org.springframework.core.MethodIntrospector;
4226
import org.springframework.core.annotation.MergedAnnotations;
43-
import org.springframework.core.env.ConfigurableEnvironment;
4427
import org.springframework.core.env.Environment;
45-
import org.springframework.javapoet.ClassName;
46-
import org.springframework.javapoet.CodeBlock;
4728
import org.springframework.test.context.DynamicPropertyRegistry;
4829
import org.springframework.test.context.DynamicPropertySource;
49-
import org.springframework.test.util.ReflectionTestUtils;
5030
import org.springframework.util.Assert;
51-
import org.springframework.util.ClassUtils;
5231
import org.springframework.util.ReflectionUtils;
5332

5433
/**
@@ -77,16 +56,6 @@ void registerDynamicPropertySources(BeanDefinitionRegistry beanDefinitionRegistr
7756
ReflectionUtils.makeAccessible(method);
7857
ReflectionUtils.invokeMethod(method, null, dynamicPropertyRegistry);
7958
});
80-
81-
String beanName = "importTestContainer.%s.%s".formatted(DynamicPropertySource.class.getName(), definitionClass);
82-
if (!beanDefinitionRegistry.containsBeanDefinition(beanName)) {
83-
RootBeanDefinition bd = new RootBeanDefinition(DynamicPropertySourceMetadata.class);
84-
bd.setInstanceSupplier(() -> new DynamicPropertySourceMetadata(definitionClass, methods));
85-
bd.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
86-
bd.setAutowireCandidate(false);
87-
bd.setAttribute(DynamicPropertySourceMetadata.class.getName(), true);
88-
beanDefinitionRegistry.registerBeanDefinition(beanName, bd);
89-
}
9059
}
9160

9261
private boolean isAnnotated(Method method) {
@@ -102,135 +71,4 @@ private void assertValid(Method method) {
10271
+ "' must accept a single DynamicPropertyRegistry argument");
10372
}
10473

105-
private record DynamicPropertySourceMetadata(Class<?> definitionClass, Set<Method> methods) {
106-
}
107-
108-
/**
109-
* {@link BeanRegistrationExcludeFilter} to exclude
110-
* {@link DynamicPropertySourceMetadata} from AOT bean registrations.
111-
*/
112-
static class DynamicPropertySourceMetadataBeanRegistrationExcludeFilter implements BeanRegistrationExcludeFilter {
113-
114-
@Override
115-
public boolean isExcludedFromAotProcessing(RegisteredBean registeredBean) {
116-
return registeredBean.getMergedBeanDefinition().hasAttribute(DynamicPropertySourceMetadata.class.getName());
117-
}
118-
119-
}
120-
121-
/**
122-
* The {@link BeanFactoryInitializationAotProcessor} generates methods for each
123-
* {@code @DynamicPropertySource-annotated} method.
124-
*
125-
*/
126-
static class DynamicPropertySourceBeanFactoryInitializationAotProcessor
127-
implements BeanFactoryInitializationAotProcessor {
128-
129-
private static final String DYNAMIC_PROPERTY_REGISTRY = "dynamicPropertyRegistry";
130-
131-
@Override
132-
public BeanFactoryInitializationAotContribution processAheadOfTime(
133-
ConfigurableListableBeanFactory beanFactory) {
134-
Map<String, DynamicPropertySourceMetadata> metadata = beanFactory
135-
.getBeansOfType(DynamicPropertySourceMetadata.class, false, false);
136-
if (metadata.isEmpty()) {
137-
return null;
138-
}
139-
return new AotContibution(metadata);
140-
}
141-
142-
private static final class AotContibution implements BeanFactoryInitializationAotContribution {
143-
144-
private final Map<String, DynamicPropertySourceMetadata> metadata;
145-
146-
private AotContibution(Map<String, DynamicPropertySourceMetadata> metadata) {
147-
this.metadata = metadata;
148-
}
149-
150-
@Override
151-
public void applyTo(GenerationContext generationContext,
152-
BeanFactoryInitializationCode beanFactoryInitializationCode) {
153-
GeneratedMethod initializerMethod = beanFactoryInitializationCode.getMethods()
154-
.add("registerDynamicPropertySources", (code) -> {
155-
code.addJavadoc("Registers {@code @DynamicPropertySource} properties");
156-
code.addParameter(ConfigurableEnvironment.class, "environment");
157-
code.addParameter(DefaultListableBeanFactory.class, "beanFactory");
158-
code.addModifiers(javax.lang.model.element.Modifier.PRIVATE,
159-
javax.lang.model.element.Modifier.STATIC);
160-
code.addStatement("$T dynamicPropertyRegistry = $T.attach(environment, beanFactory)",
161-
DynamicPropertyRegistry.class, TestcontainersPropertySource.class);
162-
this.metadata.forEach((name, metadata) -> {
163-
GeneratedMethod dynamicPropertySourceMethod = generateMethods(generationContext, metadata);
164-
code.addStatement(dynamicPropertySourceMethod.toMethodReference()
165-
.toInvokeCodeBlock(ArgumentCodeGenerator.of(DynamicPropertyRegistry.class,
166-
DYNAMIC_PROPERTY_REGISTRY)));
167-
});
168-
});
169-
beanFactoryInitializationCode.addInitializer(initializerMethod.toMethodReference());
170-
}
171-
172-
// Generates a new class in definition class package and invokes
173-
// all @DynamicPropertySource methods.
174-
private GeneratedMethod generateMethods(GenerationContext generationContext,
175-
DynamicPropertySourceMetadata metadata) {
176-
Class<?> definitionClass = metadata.definitionClass();
177-
GeneratedClass generatedClass = generationContext.getGeneratedClasses()
178-
.addForFeatureComponent(DynamicPropertySource.class.getSimpleName(), definitionClass,
179-
(code) -> code.addModifiers(javax.lang.model.element.Modifier.PUBLIC));
180-
return generatedClass.getMethods().add("registerDynamicPropertySource", (code) -> {
181-
code.addJavadoc("Registers {@code @DynamicPropertySource} properties for class '$T'",
182-
definitionClass);
183-
code.addParameter(DynamicPropertyRegistry.class, DYNAMIC_PROPERTY_REGISTRY);
184-
code.addModifiers(javax.lang.model.element.Modifier.PUBLIC,
185-
javax.lang.model.element.Modifier.STATIC);
186-
metadata.methods().forEach((method) -> {
187-
GeneratedMethod generateMethod = generateMethod(generationContext, generatedClass, method);
188-
code.addStatement(generateMethod.toMethodReference()
189-
.toInvokeCodeBlock(ArgumentCodeGenerator.of(DynamicPropertyRegistry.class,
190-
DYNAMIC_PROPERTY_REGISTRY)));
191-
});
192-
});
193-
}
194-
195-
// If the method is inaccessible, the reflection will be used; otherwise,
196-
// direct call to the method will be used.
197-
private static GeneratedMethod generateMethod(GenerationContext generationContext,
198-
GeneratedClass generatedClass, Method method) {
199-
return generatedClass.getMethods().add(method.getName(), (code) -> {
200-
code.addJavadoc("Register {@code @DynamicPropertySource} for method '$T.$L'",
201-
method.getDeclaringClass(), method.getName());
202-
code.addModifiers(javax.lang.model.element.Modifier.PRIVATE,
203-
javax.lang.model.element.Modifier.STATIC);
204-
code.addParameter(DynamicPropertyRegistry.class, DYNAMIC_PROPERTY_REGISTRY);
205-
if (isMethodAccessible(generatedClass, method)) {
206-
code.addStatement(CodeBlock.of("$T.$L($L)", method.getDeclaringClass(), method.getName(),
207-
DYNAMIC_PROPERTY_REGISTRY));
208-
}
209-
else {
210-
generationContext.getRuntimeHints().reflection().registerMethod(method, ExecutableMode.INVOKE);
211-
code.beginControlFlow("try");
212-
code.addStatement("$T<?> clazz = $T.forName($S, $T.class.getClassLoader())", Class.class,
213-
ClassUtils.class, ClassName.get(method.getDeclaringClass()), generatedClass.getName());
214-
// ReflectionTestUtils can be used here because
215-
// @DynamicPropertyRegistry in a test module.
216-
code.addStatement("$T.invokeMethod(clazz, $S, $L)", ReflectionTestUtils.class, method.getName(),
217-
DYNAMIC_PROPERTY_REGISTRY);
218-
code.nextControlFlow("catch ($T ex)", ClassNotFoundException.class);
219-
code.addStatement("throw new $T(ex)", RuntimeException.class);
220-
code.endControlFlow();
221-
}
222-
});
223-
224-
}
225-
226-
private static boolean isMethodAccessible(GeneratedClass generatedClass, Method method) {
227-
ClassName className = generatedClass.getName();
228-
return AccessControl.forClass(method.getDeclaringClass()).isAccessibleFrom(className)
229-
&& AccessControl.forMember(method).isAccessibleFrom(className);
230-
}
231-
232-
}
233-
234-
}
235-
23674
}

‎spring-boot-project/spring-boot-testcontainers/src/main/java/org/springframework/boot/testcontainers/context/ImportTestcontainersRegistrar.java

Lines changed: 132 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,34 @@
1616

1717
package org.springframework.boot.testcontainers.context;
1818

19+
import java.util.Arrays;
20+
import java.util.LinkedHashSet;
21+
import java.util.Set;
22+
23+
import javax.lang.model.element.Modifier;
24+
25+
import org.springframework.aot.AotDetector;
26+
import org.springframework.aot.generate.GeneratedClass;
27+
import org.springframework.aot.generate.GeneratedMethod;
28+
import org.springframework.aot.generate.GenerationContext;
29+
import org.springframework.aot.hint.MemberCategory;
30+
import org.springframework.aot.hint.RuntimeHints;
31+
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution;
32+
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor;
33+
import org.springframework.beans.factory.aot.BeanFactoryInitializationCode;
34+
import org.springframework.beans.factory.aot.BeanRegistrationExcludeFilter;
35+
import org.springframework.beans.factory.config.BeanDefinition;
36+
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
1937
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
38+
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
39+
import org.springframework.beans.factory.support.RegisteredBean;
40+
import org.springframework.beans.factory.support.RootBeanDefinition;
2041
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
2142
import org.springframework.core.annotation.MergedAnnotation;
43+
import org.springframework.core.env.ConfigurableEnvironment;
2244
import org.springframework.core.env.Environment;
2345
import org.springframework.core.type.AnnotationMetadata;
46+
import org.springframework.javapoet.ClassName;
2447
import org.springframework.util.ClassUtils;
2548
import org.springframework.util.ObjectUtils;
2649

@@ -51,13 +74,30 @@ public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, B
5174
MergedAnnotation<ImportTestcontainers> annotation = importingClassMetadata.getAnnotations()
5275
.get(ImportTestcontainers.class);
5376
Class<?>[] definitionClasses = annotation.getClassArray(MergedAnnotation.VALUE);
77+
Class<?> importingClass = ClassUtils.resolveClassName(importingClassMetadata.getClassName(), null);
5478
if (ObjectUtils.isEmpty(definitionClasses)) {
55-
Class<?> importingClass = ClassUtils.resolveClassName(importingClassMetadata.getClassName(), null);
5679
definitionClasses = new Class<?>[] { importingClass };
5780
}
81+
registerMetadataBeanDefinition(registry, importingClass, Set.copyOf(Arrays.asList(definitionClasses)));
5882
registerBeanDefinitions(registry, definitionClasses);
5983
}
6084

85+
private void registerMetadataBeanDefinition(BeanDefinitionRegistry registry, Class<?> importingClass,
86+
Set<Class<?>> definitionClasses) {
87+
if (!AotDetector.useGeneratedArtifacts()) {
88+
String beanName = "%s.%s.metadata".formatted(ImportTestcontainersRegistrar.class, importingClass.getName());
89+
if (registry.containsBeanDefinition(beanName)) {
90+
return;
91+
}
92+
RootBeanDefinition bd = new RootBeanDefinition(ImportTestcontainersMetadata.class);
93+
bd.setInstanceSupplier(() -> new ImportTestcontainersMetadata(importingClass, definitionClasses));
94+
bd.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
95+
bd.setAutowireCandidate(false);
96+
bd.setAttribute(ImportTestcontainersRegistrar.class.getName(), true);
97+
registry.registerBeanDefinition(beanName, bd);
98+
}
99+
}
100+
61101
private void registerBeanDefinitions(BeanDefinitionRegistry registry, Class<?>[] definitionClasses) {
62102
for (Class<?> definitionClass : definitionClasses) {
63103
this.containerFieldsImporter.registerBeanDefinitions(registry, definitionClass);
@@ -67,4 +107,95 @@ private void registerBeanDefinitions(BeanDefinitionRegistry registry, Class<?>[]
67107
}
68108
}
69109

110+
private record ImportTestcontainersMetadata(Class<?> importingClass, Set<Class<?>> definitionClasses) {
111+
}
112+
113+
static class ImportTestcontainersMetadataBeanRegistrationExcludeFilter implements BeanRegistrationExcludeFilter {
114+
115+
@Override
116+
public boolean isExcludedFromAotProcessing(RegisteredBean registeredBean) {
117+
return registeredBean.getMergedBeanDefinition().hasAttribute(ImportTestcontainersRegistrar.class.getName());
118+
}
119+
120+
}
121+
122+
static class ImportTestcontainersBeanFactoryInitializationAotProcessor
123+
implements BeanFactoryInitializationAotProcessor {
124+
125+
@Override
126+
public BeanFactoryInitializationAotContribution processAheadOfTime(
127+
ConfigurableListableBeanFactory beanFactory) {
128+
Set<ImportTestcontainersMetadata> importClasses = new LinkedHashSet<>(
129+
beanFactory.getBeansOfType(ImportTestcontainersMetadata.class, false, false).values());
130+
if (importClasses.isEmpty()) {
131+
return null;
132+
}
133+
return new AotContibution(importClasses);
134+
}
135+
136+
private static final class AotContibution implements BeanFactoryInitializationAotContribution {
137+
138+
private static final String BEAN_FACTORY_PARAM = "beanFactory";
139+
140+
private static final String ENVIRONMENT_PARAM = "environment";
141+
142+
private static final String IMPORTING_CLASS_PARAM = "importingClass";
143+
144+
private final Set<ImportTestcontainersMetadata> metadata;
145+
146+
private AotContibution(Set<ImportTestcontainersMetadata> metadata) {
147+
this.metadata = metadata;
148+
}
149+
150+
@Override
151+
public void applyTo(GenerationContext generationContext,
152+
BeanFactoryInitializationCode beanFactoryInitializationCode) {
153+
154+
contributeHints(generationContext.getRuntimeHints());
155+
156+
GeneratedClass generatedClass = generationContext.getGeneratedClasses()
157+
.addForFeatureComponent(ImportTestcontainers.class.getSimpleName(),
158+
ImportTestcontainersRegistrar.class, (code) -> code.addModifiers(Modifier.PUBLIC));
159+
160+
GeneratedMethod importBeanDefinitionMethod = generateImportBeanDefinitionMethod(generatedClass);
161+
GeneratedMethod initializeMethod = generatedClass.getMethods()
162+
.add("registerBeanDefinitions", (code) -> {
163+
code.addJavadoc("Register bean definitions for '$T'", ImportTestcontainers.class);
164+
code.addModifiers(Modifier.PUBLIC, Modifier.STATIC);
165+
code.addParameter(ConfigurableEnvironment.class, ENVIRONMENT_PARAM);
166+
code.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAM);
167+
this.metadata.forEach((metadata) -> code.addStatement("$L($L, $L, $S)",
168+
importBeanDefinitionMethod.getName(), ENVIRONMENT_PARAM, BEAN_FACTORY_PARAM,
169+
ClassName.get(metadata.importingClass())));
170+
});
171+
beanFactoryInitializationCode.addInitializer(initializeMethod.toMethodReference());
172+
}
173+
174+
private void contributeHints(RuntimeHints runtimeHints) {
175+
Set<Class<?>> definitionClasses = new LinkedHashSet<>();
176+
this.metadata.forEach((metadata) -> definitionClasses.addAll(metadata.definitionClasses()));
177+
definitionClasses.forEach((definitionClass) -> runtimeHints.reflection()
178+
.registerType(definitionClass, MemberCategory.DECLARED_FIELDS, MemberCategory.PUBLIC_FIELDS,
179+
MemberCategory.INVOKE_PUBLIC_METHODS, MemberCategory.INVOKE_DECLARED_METHODS));
180+
}
181+
182+
private GeneratedMethod generateImportBeanDefinitionMethod(GeneratedClass generatedClass) {
183+
return generatedClass.getMethods().add("registerBeanDefinitionsFor", (code) -> {
184+
code.addModifiers(Modifier.PRIVATE, Modifier.STATIC);
185+
code.addParameter(ConfigurableEnvironment.class, ENVIRONMENT_PARAM);
186+
code.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAM);
187+
code.addParameter(String.class, IMPORTING_CLASS_PARAM);
188+
code.addStatement("$T<?> clazz = $T.resolveClassName($L, $L.getBeanClassLoader())", Class.class,
189+
ClassUtils.class, IMPORTING_CLASS_PARAM, BEAN_FACTORY_PARAM);
190+
code.addStatement("$T metadata = $T.introspect(clazz)", AnnotationMetadata.class,
191+
AnnotationMetadata.class);
192+
code.addStatement("new $T($L).registerBeanDefinitions(metadata, $L)",
193+
ImportTestcontainersRegistrar.class, ENVIRONMENT_PARAM, BEAN_FACTORY_PARAM);
194+
});
195+
}
196+
197+
}
198+
199+
}
200+
70201
}

0 commit comments

Comments
 (0)