Skip to content

Commit cd63de3

Browse files
nosannosan
committed
TestcontainersBeanRegistrationAotProcessor that replaces InstanceSupplier of Container by either direct field usage or a reflection equivalent.
If the field is private, the reflection will be used; otherwise, direct access to the field will be used DynamicPropertySourceBeanFactoryInitializationAotProcessor that generates methods for each annotated @DynamicPropertySource method
1 parent 4718485 commit cd63de3

File tree

5 files changed

+382
-3
lines changed

5 files changed

+382
-3
lines changed

‎spring-boot-project/spring-boot-testcontainers/src/dockerTest/java/org/springframework/boot/testcontainers/ImportTestcontainersTests.java

+120
Original file line numberDiff line numberDiff line change
@@ -18,17 +18,27 @@
1818

1919
import java.lang.annotation.Retention;
2020
import java.lang.annotation.RetentionPolicy;
21+
import java.util.function.BiConsumer;
2122

2223
import org.junit.jupiter.api.AfterEach;
2324
import org.junit.jupiter.api.Test;
2425
import org.testcontainers.containers.Container;
2526
import org.testcontainers.containers.PostgreSQLContainer;
2627

28+
import org.springframework.aot.test.generate.TestGenerationContext;
2729
import org.springframework.boot.testcontainers.beans.TestcontainerBeanDefinition;
2830
import org.springframework.boot.testcontainers.context.ImportTestcontainers;
31+
import org.springframework.boot.testcontainers.lifecycle.TestcontainersLifecycleApplicationContextInitializer;
2932
import org.springframework.boot.testsupport.container.DisabledIfDockerUnavailable;
3033
import org.springframework.boot.testsupport.container.TestImage;
34+
import org.springframework.context.ApplicationContextInitializer;
3135
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
36+
import org.springframework.context.aot.ApplicationContextAotGenerator;
37+
import org.springframework.context.support.GenericApplicationContext;
38+
import org.springframework.core.test.tools.CompileWithForkedClassLoader;
39+
import org.springframework.core.test.tools.Compiled;
40+
import org.springframework.core.test.tools.TestCompiler;
41+
import org.springframework.javapoet.ClassName;
3242
import org.springframework.test.context.DynamicPropertyRegistry;
3343
import org.springframework.test.context.DynamicPropertySource;
3444

@@ -43,6 +53,8 @@
4353
@DisabledIfDockerUnavailable
4454
class ImportTestcontainersTests {
4555

56+
private final TestGenerationContext generationContext = new TestGenerationContext();
57+
4658
private AnnotationConfigApplicationContext applicationContext;
4759

4860
@AfterEach
@@ -122,6 +134,87 @@ void importWhenHasBadArgsDynamicPropertySourceMethod() {
122134
.withMessage("@DynamicPropertySource method 'containerProperties' must be static");
123135
}
124136

137+
@Test
138+
@CompileWithForkedClassLoader
139+
void importTestcontainersImportWithoutValueAotContribution() {
140+
this.applicationContext = new AnnotationConfigApplicationContext();
141+
this.applicationContext.register(ImportWithoutValue.class);
142+
compile((freshContext, compiled) -> {
143+
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
144+
assertThat(container).isSameAs(ImportWithoutValue.container);
145+
});
146+
}
147+
148+
@Test
149+
@CompileWithForkedClassLoader
150+
void importTestcontainersImportWithValueAotContribution() {
151+
this.applicationContext = new AnnotationConfigApplicationContext();
152+
this.applicationContext.register(ImportWithValue.class);
153+
compile((freshContext, compiled) -> {
154+
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
155+
assertThat(container).isSameAs(ContainerDefinitions.container);
156+
});
157+
}
158+
159+
@Test
160+
@CompileWithForkedClassLoader
161+
void importTestcontainersWithDynamicPropertySourceAotContribution() {
162+
this.applicationContext = new AnnotationConfigApplicationContext();
163+
this.applicationContext.register(ContainerDefinitionsWithDynamicPropertySource.class);
164+
compile((freshContext, compiled) -> {
165+
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
166+
assertThat(container).isSameAs(ContainerDefinitionsWithDynamicPropertySource.container);
167+
assertThat(freshContext.getEnvironment().getProperty("container.port", Integer.class))
168+
.isEqualTo(ContainerDefinitionsWithDynamicPropertySource.container.getFirstMappedPort());
169+
});
170+
}
171+
172+
@Test
173+
@CompileWithForkedClassLoader
174+
void importTestcontainersWithCustomPostgreSQLContainerAotContribution() {
175+
this.applicationContext = new AnnotationConfigApplicationContext();
176+
this.applicationContext.register(CustomPostgreSQLContainerDefinitions.class);
177+
compile((freshContext, compiled) -> {
178+
CustomPostgreSQLContainer container = freshContext.getBean(CustomPostgreSQLContainer.class);
179+
assertThat(container).isSameAs(CustomPostgreSQLContainerDefinitions.container);
180+
});
181+
}
182+
183+
@Test
184+
@CompileWithForkedClassLoader
185+
void importTestcontainersWithNotAccessibleContainerAndDynamicPropertySourceAotContribution() {
186+
this.applicationContext = new AnnotationConfigApplicationContext();
187+
this.applicationContext.register(ImportNotAccessibleContainerAndDynamicPropertySource.class);
188+
compile((freshContext, compiled) -> {
189+
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
190+
assertThat(container).isSameAs(ImportNotAccessibleContainerAndDynamicPropertySource.container);
191+
assertThat(freshContext.getEnvironment().getProperty("container.port", Integer.class))
192+
.isEqualTo(ImportNotAccessibleContainerAndDynamicPropertySource.container.getFirstMappedPort());
193+
});
194+
}
195+
196+
@SuppressWarnings("unchecked")
197+
private void compile(BiConsumer<GenericApplicationContext, Compiled> result) {
198+
ClassName className = processAheadOfTime();
199+
TestCompiler.forSystem().with(this.generationContext).compile((compiled) -> {
200+
try (GenericApplicationContext context = new GenericApplicationContext()) {
201+
new TestcontainersLifecycleApplicationContextInitializer().initialize(context);
202+
ApplicationContextInitializer<GenericApplicationContext> initializer = compiled
203+
.getInstance(ApplicationContextInitializer.class, className.toString());
204+
initializer.initialize(context);
205+
context.refresh();
206+
result.accept(context, compiled);
207+
}
208+
});
209+
}
210+
211+
private ClassName processAheadOfTime() {
212+
ClassName className = new ApplicationContextAotGenerator().processAheadOfTime(this.applicationContext,
213+
this.generationContext);
214+
this.generationContext.writeGeneratedContent();
215+
return className;
216+
}
217+
125218
@ImportTestcontainers
126219
static class ImportWithoutValue {
127220

@@ -196,4 +289,31 @@ void containerProperties() {
196289

197290
}
198291

292+
@ImportTestcontainers
293+
static class CustomPostgreSQLContainerDefinitions {
294+
295+
static CustomPostgreSQLContainer container = new CustomPostgreSQLContainer();
296+
297+
}
298+
299+
static class CustomPostgreSQLContainer extends PostgreSQLContainer<CustomPostgreSQLContainer> {
300+
301+
CustomPostgreSQLContainer() {
302+
super("postgres:14");
303+
}
304+
305+
}
306+
307+
@ImportTestcontainers
308+
static class ImportNotAccessibleContainerAndDynamicPropertySource {
309+
310+
private static final PostgreSQLContainer<?> container = TestImage.container(PostgreSQLContainer.class);
311+
312+
@DynamicPropertySource
313+
private static void containerProperties(DynamicPropertyRegistry registry) {
314+
registry.add("container.port", container::getFirstMappedPort);
315+
}
316+
317+
}
318+
199319
}

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

+150
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,31 @@
1818

1919
import java.lang.reflect.Method;
2020
import java.lang.reflect.Modifier;
21+
import java.util.Map;
2122
import java.util.Set;
2223

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;
2336
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;
2440
import org.springframework.boot.testcontainers.properties.TestcontainersPropertySource;
2541
import org.springframework.core.MethodIntrospector;
2642
import org.springframework.core.annotation.MergedAnnotations;
43+
import org.springframework.core.env.ConfigurableEnvironment;
2744
import org.springframework.core.env.Environment;
45+
import org.springframework.javapoet.CodeBlock;
2846
import org.springframework.test.context.DynamicPropertyRegistry;
2947
import org.springframework.test.context.DynamicPropertySource;
3048
import org.springframework.util.Assert;
@@ -56,6 +74,16 @@ void registerDynamicPropertySources(BeanDefinitionRegistry beanDefinitionRegistr
5674
ReflectionUtils.makeAccessible(method);
5775
ReflectionUtils.invokeMethod(method, null, dynamicPropertyRegistry);
5876
});
77+
78+
String beanName = "importTestContainer.%s.%s".formatted(DynamicPropertySource.class.getName(), definitionClass);
79+
if (!beanDefinitionRegistry.containsBeanDefinition(beanName)) {
80+
RootBeanDefinition bd = new RootBeanDefinition(DynamicPropertySourceMetadata.class);
81+
bd.setInstanceSupplier(() -> new DynamicPropertySourceMetadata(definitionClass, methods));
82+
bd.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
83+
bd.setAutowireCandidate(false);
84+
bd.setAttribute(DynamicPropertySourceMetadata.class.getName(), true);
85+
beanDefinitionRegistry.registerBeanDefinition(beanName, bd);
86+
}
5987
}
6088

6189
private boolean isAnnotated(Method method) {
@@ -71,4 +99,126 @@ private void assertValid(Method method) {
7199
+ "' must accept a single DynamicPropertyRegistry argument");
72100
}
73101

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

0 commit comments

Comments
 (0)