Skip to content

Commit 7c38b50

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 7c38b50

File tree

5 files changed

+131
-276
lines changed

5 files changed

+131
-276
lines changed

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

-162
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

+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)