Skip to content

Commit 1cc38b6

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 682436e commit 1cc38b6

File tree

5 files changed

+131
-272
lines changed

5 files changed

+131
-272
lines changed

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

-158
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,15 +56,6 @@ void registerDynamicPropertySources(BeanDefinitionRegistry beanDefinitionRegistr
7756
ReflectionUtils.makeAccessible(method);
7857
ReflectionUtils.invokeMethod(method, null, dynamicPropertyRegistry);
7958
});
80-
String beanName = "%s.%s".formatted(DynamicPropertySourceMetadata.class.getName(), definitionClass);
81-
if (!beanDefinitionRegistry.containsBeanDefinition(beanName)) {
82-
RootBeanDefinition bd = new RootBeanDefinition(DynamicPropertySourceMetadata.class);
83-
bd.setInstanceSupplier(() -> new DynamicPropertySourceMetadata(definitionClass, methods));
84-
bd.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
85-
bd.setAutowireCandidate(false);
86-
bd.setAttribute(DynamicPropertySourceMetadata.class.getName(), true);
87-
beanDefinitionRegistry.registerBeanDefinition(beanName, bd);
88-
}
8959
}
9060

9161
private boolean isAnnotated(Method method) {
@@ -101,132 +71,4 @@ private void assertValid(Method method) {
10171
+ "' must accept a single DynamicPropertyRegistry argument");
10272
}
10373

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

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

+128-2
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,32 @@
1616

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

19+
import java.util.LinkedHashSet;
20+
import java.util.Map;
21+
import java.util.Set;
22+
import java.util.stream.Collectors;
23+
import java.util.stream.Stream;
24+
25+
import javax.lang.model.element.Modifier;
26+
27+
import org.springframework.aot.generate.GeneratedClass;
28+
import org.springframework.aot.generate.GeneratedMethod;
29+
import org.springframework.aot.generate.GenerationContext;
30+
import org.springframework.aot.hint.MemberCategory;
31+
import org.springframework.aot.hint.RuntimeHints;
32+
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution;
33+
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor;
34+
import org.springframework.beans.factory.aot.BeanFactoryInitializationCode;
35+
import org.springframework.beans.factory.aot.BeanRegistrationExcludeFilter;
36+
import org.springframework.beans.factory.config.BeanDefinition;
37+
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
1938
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
39+
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
40+
import org.springframework.beans.factory.support.RegisteredBean;
41+
import org.springframework.beans.factory.support.RootBeanDefinition;
2042
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
2143
import org.springframework.core.annotation.MergedAnnotation;
44+
import org.springframework.core.env.ConfigurableEnvironment;
2245
import org.springframework.core.env.Environment;
2346
import org.springframework.core.type.AnnotationMetadata;
2447
import org.springframework.util.ClassUtils;
@@ -51,14 +74,15 @@ 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, definitionClasses);
5882
registerBeanDefinitions(registry, definitionClasses);
5983
}
6084

61-
private void registerBeanDefinitions(BeanDefinitionRegistry registry, Class<?>[] definitionClasses) {
85+
void registerBeanDefinitions(BeanDefinitionRegistry registry, Class<?>[] definitionClasses) {
6286
for (Class<?> definitionClass : definitionClasses) {
6387
this.containerFieldsImporter.registerBeanDefinitions(registry, definitionClass);
6488
if (this.dynamicPropertySourceMethodsImporter != null) {
@@ -67,4 +91,106 @@ private void registerBeanDefinitions(BeanDefinitionRegistry registry, Class<?>[]
6791
}
6892
}
6993

94+
private void registerMetadataBeanDefinition(BeanDefinitionRegistry registry, Class<?> importingClass,
95+
Class<?>[] definitionClasses) {
96+
String beanName = "%s.%s.metadata".formatted(ImportTestcontainersMetadata.class, importingClass.getName());
97+
if (!registry.containsBeanDefinition(beanName)) {
98+
RootBeanDefinition bd = new RootBeanDefinition(ImportTestcontainersMetadata.class);
99+
bd.setInstanceSupplier(() -> new ImportTestcontainersMetadata(importingClass, definitionClasses));
100+
bd.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
101+
bd.setAutowireCandidate(false);
102+
bd.setAttribute(ImportTestcontainersMetadata.class.getName(), true);
103+
registry.registerBeanDefinition(beanName, bd);
104+
}
105+
}
106+
107+
private record ImportTestcontainersMetadata(Class<?> importingClass, Class<?>[] definitionClasses) {
108+
}
109+
110+
static class ImportTestcontainersBeanRegistrationExcludeFilter implements BeanRegistrationExcludeFilter {
111+
112+
@Override
113+
public boolean isExcludedFromAotProcessing(RegisteredBean registeredBean) {
114+
RootBeanDefinition bd = registeredBean.getMergedBeanDefinition();
115+
return bd.hasAttribute(TestcontainerFieldBeanDefinition.class.getName())
116+
|| bd.hasAttribute(ImportTestcontainersMetadata.class.getName());
117+
}
118+
119+
}
120+
121+
static class ImportTestcontainersBeanFactoryInitializationAotProcessor
122+
implements BeanFactoryInitializationAotProcessor {
123+
124+
@Override
125+
public BeanFactoryInitializationAotContribution processAheadOfTime(
126+
ConfigurableListableBeanFactory beanFactory) {
127+
Map<String, ImportTestcontainersMetadata> metadata = beanFactory
128+
.getBeansOfType(ImportTestcontainersMetadata.class, false, false);
129+
if (metadata.isEmpty()) {
130+
return null;
131+
}
132+
return new AotContribution(new LinkedHashSet<>(metadata.values()));
133+
}
134+
135+
private static final class AotContribution implements BeanFactoryInitializationAotContribution {
136+
137+
private static final String BEAN_FACTORY_PARAM = "beanFactory";
138+
139+
private static final String ENVIRONMENT_PARAM = "environment";
140+
141+
private final Set<ImportTestcontainersMetadata> metadata;
142+
143+
private AotContribution(Set<ImportTestcontainersMetadata> metadata) {
144+
this.metadata = metadata;
145+
}
146+
147+
@Override
148+
public void applyTo(GenerationContext generationContext,
149+
BeanFactoryInitializationCode beanFactoryInitializationCode) {
150+
151+
Set<Class<?>> definitionClasses = getDefinitionClasses();
152+
contributeHints(generationContext.getRuntimeHints(), definitionClasses);
153+
154+
GeneratedClass generatedClass = generationContext.getGeneratedClasses()
155+
.addForFeatureComponent(ImportTestcontainers.class.getSimpleName(),
156+
ImportTestcontainersRegistrar.class, (code) -> code.addModifiers(Modifier.PUBLIC));
157+
158+
GeneratedMethod initializeMethod = generatedClass.getMethods()
159+
.add("registerBeanDefinitions", (code) -> {
160+
code.addJavadoc("Register bean definitions for '$T'", ImportTestcontainers.class);
161+
code.addModifiers(Modifier.PUBLIC, Modifier.STATIC);
162+
code.addParameter(ConfigurableEnvironment.class, ENVIRONMENT_PARAM);
163+
code.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAM);
164+
code.addStatement("$T<$T<?>> definitionClasses = new $T<>()", Set.class, Class.class,
165+
LinkedHashSet.class);
166+
code.addStatement("$T classLoader = $L.getBeanClassLoader()", ClassLoader.class,
167+
BEAN_FACTORY_PARAM);
168+
definitionClasses.forEach((definitionClass) -> code.addStatement(
169+
"definitionClasses.add($T.resolveClassName($S, classLoader))", ClassUtils.class,
170+
definitionClass.getTypeName()));
171+
code.addStatement(
172+
"new $T($L).registerBeanDefinitions($L, definitionClasses.toArray(new $T<?>[0]))",
173+
ImportTestcontainersRegistrar.class, ENVIRONMENT_PARAM, BEAN_FACTORY_PARAM,
174+
Class.class);
175+
});
176+
beanFactoryInitializationCode.addInitializer(initializeMethod.toMethodReference());
177+
}
178+
179+
private Set<Class<?>> getDefinitionClasses() {
180+
return this.metadata.stream()
181+
.map(ImportTestcontainersMetadata::definitionClasses)
182+
.flatMap(Stream::of)
183+
.collect(Collectors.toCollection(LinkedHashSet::new));
184+
}
185+
186+
private void contributeHints(RuntimeHints runtimeHints, Set<Class<?>> definitionClasses) {
187+
definitionClasses.forEach((definitionClass) -> runtimeHints.reflection()
188+
.registerType(definitionClass, MemberCategory.DECLARED_FIELDS, MemberCategory.PUBLIC_FIELDS,
189+
MemberCategory.INVOKE_PUBLIC_METHODS, MemberCategory.INVOKE_DECLARED_METHODS));
190+
}
191+
192+
}
193+
194+
}
195+
70196
}

0 commit comments

Comments
 (0)