Skip to content

Commit cb9f9b5

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
1 parent 4718485 commit cb9f9b5

File tree

5 files changed

+334
-4
lines changed

5 files changed

+334
-4
lines changed

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

+112
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,84 @@ 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+
new TestcontainersLifecycleApplicationContextInitializer().initialize(this.applicationContext);
165+
compile((freshContext, compiled) -> {
166+
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
167+
assertThat(container).isSameAs(ContainerDefinitionsWithDynamicPropertySource.container);
168+
assertThat(freshContext.getEnvironment().getProperty("container.port", Integer.class))
169+
.isEqualTo(ContainerDefinitionsWithDynamicPropertySource.container.getFirstMappedPort());
170+
});
171+
}
172+
173+
@Test
174+
@CompileWithForkedClassLoader
175+
void importTestcontainersWithCustomPostgreSQLContainerAotContribution() {
176+
this.applicationContext = new AnnotationConfigApplicationContext();
177+
this.applicationContext.register(CustomPostgreSQLContainerDefinitions.class);
178+
compile((freshContext, compiled) -> {
179+
CustomPostgreSQLContainer container = freshContext.getBean(CustomPostgreSQLContainer.class);
180+
assertThat(container).isSameAs(CustomPostgreSQLContainerDefinitions.container);
181+
});
182+
}
183+
184+
@Test
185+
@CompileWithForkedClassLoader
186+
void importTestcontainersWithNotAccessibleContainerAotContribution() {
187+
this.applicationContext = new AnnotationConfigApplicationContext();
188+
this.applicationContext.register(ImportNotAccessibleContainer.class);
189+
compile((freshContext, compiled) -> {
190+
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
191+
assertThat(container).isSameAs(ImportNotAccessibleContainer.container);
192+
});
193+
}
194+
195+
@SuppressWarnings("unchecked")
196+
private void compile(BiConsumer<GenericApplicationContext, Compiled> result) {
197+
ClassName className = processAheadOfTime();
198+
TestCompiler.forSystem().with(this.generationContext).compile((compiled) -> {
199+
GenericApplicationContext freshApplicationContext = new GenericApplicationContext();
200+
ApplicationContextInitializer<GenericApplicationContext> initializer = compiled
201+
.getInstance(ApplicationContextInitializer.class, className.toString());
202+
initializer.initialize(freshApplicationContext);
203+
freshApplicationContext.refresh();
204+
result.accept(freshApplicationContext, compiled);
205+
});
206+
}
207+
208+
private ClassName processAheadOfTime() {
209+
ClassName className = new ApplicationContextAotGenerator().processAheadOfTime(this.applicationContext,
210+
this.generationContext);
211+
this.generationContext.writeGeneratedContent();
212+
return className;
213+
}
214+
125215
@ImportTestcontainers
126216
static class ImportWithoutValue {
127217

@@ -196,4 +286,26 @@ void containerProperties() {
196286

197287
}
198288

289+
@ImportTestcontainers
290+
static class CustomPostgreSQLContainerDefinitions {
291+
292+
static CustomPostgreSQLContainer container = new CustomPostgreSQLContainer();
293+
294+
}
295+
296+
static class CustomPostgreSQLContainer extends PostgreSQLContainer<CustomPostgreSQLContainer> {
297+
298+
CustomPostgreSQLContainer() {
299+
super("postgres:14");
300+
}
301+
302+
}
303+
304+
@ImportTestcontainers
305+
static class ImportNotAccessibleContainer {
306+
307+
private static final PostgreSQLContainer<?> container = TestImage.container(PostgreSQLContainer.class);
308+
309+
}
310+
199311
}

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

+109
Original file line numberDiff line numberDiff line change
@@ -18,16 +18,32 @@
1818

1919
import java.lang.reflect.Method;
2020
import java.lang.reflect.Modifier;
21+
import java.util.LinkedHashSet;
22+
import java.util.Map;
2123
import java.util.Set;
2224

25+
import org.springframework.aot.generate.GeneratedMethod;
26+
import org.springframework.aot.generate.GenerationContext;
27+
import org.springframework.aot.hint.ExecutableMode;
28+
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution;
29+
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor;
30+
import org.springframework.beans.factory.aot.BeanFactoryInitializationCode;
31+
import org.springframework.beans.factory.aot.BeanRegistrationExcludeFilter;
32+
import org.springframework.beans.factory.config.BeanDefinition;
33+
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
2334
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
35+
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
36+
import org.springframework.beans.factory.support.RegisteredBean;
37+
import org.springframework.beans.factory.support.RootBeanDefinition;
2438
import org.springframework.boot.testcontainers.properties.TestcontainersPropertySource;
2539
import org.springframework.core.MethodIntrospector;
2640
import org.springframework.core.annotation.MergedAnnotations;
41+
import org.springframework.core.env.ConfigurableEnvironment;
2742
import org.springframework.core.env.Environment;
2843
import org.springframework.test.context.DynamicPropertyRegistry;
2944
import org.springframework.test.context.DynamicPropertySource;
3045
import org.springframework.util.Assert;
46+
import org.springframework.util.ClassUtils;
3147
import org.springframework.util.ReflectionUtils;
3248

3349
/**
@@ -51,11 +67,22 @@ void registerDynamicPropertySources(BeanDefinitionRegistry beanDefinitionRegistr
5167
}
5268
DynamicPropertyRegistry dynamicPropertyRegistry = TestcontainersPropertySource.attach(this.environment,
5369
beanDefinitionRegistry);
70+
DynamicPropertySourceMethodsImporterMetadata metadata = new DynamicPropertySourceMethodsImporterMetadata();
5471
methods.forEach((method) -> {
5572
assertValid(method);
5673
ReflectionUtils.makeAccessible(method);
5774
ReflectionUtils.invokeMethod(method, null, dynamicPropertyRegistry);
75+
metadata.methods.add(method);
5876
});
77+
String beanName = "importTestContainer.%s.%s".formatted(DynamicPropertySource.class.getName(), definitionClass);
78+
if (!beanDefinitionRegistry.containsBeanDefinition(beanName)) {
79+
RootBeanDefinition bd = new RootBeanDefinition(DynamicPropertySourceMethodsImporterMetadata.class);
80+
bd.setInstanceSupplier(() -> metadata);
81+
bd.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
82+
bd.setAutowireCandidate(false);
83+
bd.setAttribute(DynamicPropertySourceMethodsImporterMetadata.class.getName(), true);
84+
beanDefinitionRegistry.registerBeanDefinition(beanName, bd);
85+
}
5986
}
6087

6188
private boolean isAnnotated(Method method) {
@@ -71,4 +98,86 @@ private void assertValid(Method method) {
7198
+ "' must accept a single DynamicPropertyRegistry argument");
7299
}
73100

101+
private static final class DynamicPropertySourceMethodsImporterMetadata {
102+
103+
private final Set<Method> methods = new LinkedHashSet<>();
104+
105+
}
106+
107+
static class DynamicPropertySourceMethodsImporterMetadataBeanRegistrationExcludeFilter
108+
implements BeanRegistrationExcludeFilter {
109+
110+
@Override
111+
public boolean isExcludedFromAotProcessing(RegisteredBean registeredBean) {
112+
return registeredBean.getMergedBeanDefinition()
113+
.hasAttribute(DynamicPropertySourceMethodsImporterMetadata.class.getName());
114+
}
115+
116+
}
117+
118+
/**
119+
* {@link BeanFactoryInitializationAotProcessor} that generates all
120+
* {@link DynamicPropertySource} methods if any.
121+
*
122+
*/
123+
static class DynamicPropertySourceBeanFactoryInitializationAotProcessor
124+
implements BeanFactoryInitializationAotProcessor {
125+
126+
@Override
127+
public BeanFactoryInitializationAotContribution processAheadOfTime(
128+
ConfigurableListableBeanFactory beanFactory) {
129+
Map<String, DynamicPropertySourceMethodsImporterMetadata> metadata = beanFactory
130+
.getBeansOfType(DynamicPropertySourceMethodsImporterMetadata.class);
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, DynamicPropertySourceMethodsImporterMetadata> metadata;
140+
141+
private AotContibution(Map<String, DynamicPropertySourceMethodsImporterMetadata> metadata) {
142+
this.metadata = metadata;
143+
}
144+
145+
@Override
146+
public void applyTo(GenerationContext generationContext,
147+
BeanFactoryInitializationCode beanFactoryInitializationCode) {
148+
this.metadata.forEach((name, metadata) -> metadata.methods.forEach((method) -> {
149+
generationContext.getRuntimeHints().reflection().registerMethod(method, ExecutableMode.INVOKE);
150+
GeneratedMethod generatedMethod = beanFactoryInitializationCode.getMethods()
151+
.add(method.getName(), (code) -> {
152+
code.addJavadoc("DynamicPropertySource for method $L.$L",
153+
method.getDeclaringClass().getName(), method.getName());
154+
code.addModifiers(javax.lang.model.element.Modifier.PRIVATE,
155+
javax.lang.model.element.Modifier.STATIC);
156+
code.addParameter(ConfigurableEnvironment.class, "environment");
157+
code.addParameter(DefaultListableBeanFactory.class, "beanFactory");
158+
code.addStatement("$T dynamicPropertyRegistry = $T.attach(environment, beanFactory)",
159+
DynamicPropertyRegistry.class, TestcontainersPropertySource.class);
160+
code.beginControlFlow("try");
161+
code.addStatement("$T<?> clazz = $T.forName($S, beanFactory.getBeanClassLoader())",
162+
Class.class, ClassUtils.class, method.getDeclaringClass().getName());
163+
code.addStatement("$T method = $T.findMethod(clazz, $S, $T.class)", Method.class,
164+
ReflectionUtils.class, method.getName(), DynamicPropertyRegistry.class);
165+
code.addStatement("$T.notNull(method, $S)", Assert.class,
166+
"Method '" + method.getName() + "' is not found");
167+
code.addStatement("$T.makeAccessible(method)", ReflectionUtils.class);
168+
code.addStatement("$T.invokeMethod(method, null, dynamicPropertyRegistry)",
169+
ReflectionUtils.class);
170+
code.nextControlFlow("catch ($T ex)", ClassNotFoundException.class);
171+
code.addStatement("throw new $T(ex)", RuntimeException.class);
172+
code.endControlFlow();
173+
});
174+
beanFactoryInitializationCode.addInitializer(generatedMethod.toMethodReference());
175+
}));
176+
177+
}
178+
179+
}
180+
181+
}
182+
74183
}

0 commit comments

Comments
 (0)