Skip to content

Commit 91e3273

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 91e3273

File tree

5 files changed

+373
-3
lines changed

5 files changed

+373
-3
lines changed

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

Lines changed: 120 additions & 0 deletions
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 importTestcontainersWithNotAccessibleContainerAotContribution() {
186+
this.applicationContext = new AnnotationConfigApplicationContext();
187+
this.applicationContext.register(ImportNotAccessibleContainer.class);
188+
compile((freshContext, compiled) -> {
189+
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
190+
assertThat(container).isSameAs(ImportNotAccessibleContainer.container);
191+
assertThat(freshContext.getEnvironment().getProperty("container.port", Integer.class))
192+
.isEqualTo(ImportNotAccessibleContainer.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 ImportNotAccessibleContainer {
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‎

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

0 commit comments

Comments
 (0)