Skip to content

Commit

Permalink
TestcontainersBeanRegistrationAotProcessor that replaces InstanceSupp…
Browse files Browse the repository at this point in the history 
…lier 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
  • Loading branch information
@nosan
nosan committed Oct 27, 2024
1 parent 4718485 commit 91e3273
Show file tree
Hide file tree
Showing 5 changed files with 373 additions and 3 deletions.
120 changes: 120 additions & 0 deletions ...rc/dockerTest/java/org/springframework/boot/testcontainers/ImportTestcontainersTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,27 @@

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.util.function.BiConsumer;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.Container;
import org.testcontainers.containers.PostgreSQLContainer;

import org.springframework.aot.test.generate.TestGenerationContext;
import org.springframework.boot.testcontainers.beans.TestcontainerBeanDefinition;
import org.springframework.boot.testcontainers.context.ImportTestcontainers;
import org.springframework.boot.testcontainers.lifecycle.TestcontainersLifecycleApplicationContextInitializer;
import org.springframework.boot.testsupport.container.DisabledIfDockerUnavailable;
import org.springframework.boot.testsupport.container.TestImage;
import org.springframework.context.ApplicationContextInitializer;
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
import org.springframework.context.aot.ApplicationContextAotGenerator;
import org.springframework.context.support.GenericApplicationContext;
import org.springframework.core.test.tools.CompileWithForkedClassLoader;
import org.springframework.core.test.tools.Compiled;
import org.springframework.core.test.tools.TestCompiler;
import org.springframework.javapoet.ClassName;
import org.springframework.test.context.DynamicPropertyRegistry;
import org.springframework.test.context.DynamicPropertySource;

Expand All @@ -43,6 +53,8 @@
@DisabledIfDockerUnavailable
class ImportTestcontainersTests {

private final TestGenerationContext generationContext = new TestGenerationContext();

private AnnotationConfigApplicationContext applicationContext;

@AfterEach
Expand Down Expand Up @@ -122,6 +134,87 @@ void importWhenHasBadArgsDynamicPropertySourceMethod() {
.withMessage("@DynamicPropertySource method 'containerProperties' must be static");
}

@Test
@CompileWithForkedClassLoader
void importTestcontainersImportWithoutValueAotContribution() {
this.applicationContext = new AnnotationConfigApplicationContext();
this.applicationContext.register(ImportWithoutValue.class);
compile((freshContext, compiled) -> {
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
assertThat(container).isSameAs(ImportWithoutValue.container);
});
}

@Test
@CompileWithForkedClassLoader
void importTestcontainersImportWithValueAotContribution() {
this.applicationContext = new AnnotationConfigApplicationContext();
this.applicationContext.register(ImportWithValue.class);
compile((freshContext, compiled) -> {
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
assertThat(container).isSameAs(ContainerDefinitions.container);
});
}

@Test
@CompileWithForkedClassLoader
void importTestcontainersWithDynamicPropertySourceAotContribution() {
this.applicationContext = new AnnotationConfigApplicationContext();
this.applicationContext.register(ContainerDefinitionsWithDynamicPropertySource.class);
compile((freshContext, compiled) -> {
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
assertThat(container).isSameAs(ContainerDefinitionsWithDynamicPropertySource.container);
assertThat(freshContext.getEnvironment().getProperty("container.port", Integer.class))
.isEqualTo(ContainerDefinitionsWithDynamicPropertySource.container.getFirstMappedPort());
});
}

@Test
@CompileWithForkedClassLoader
void importTestcontainersWithCustomPostgreSQLContainerAotContribution() {
this.applicationContext = new AnnotationConfigApplicationContext();
this.applicationContext.register(CustomPostgreSQLContainerDefinitions.class);
compile((freshContext, compiled) -> {
CustomPostgreSQLContainer container = freshContext.getBean(CustomPostgreSQLContainer.class);
assertThat(container).isSameAs(CustomPostgreSQLContainerDefinitions.container);
});
}

@Test
@CompileWithForkedClassLoader
void importTestcontainersWithNotAccessibleContainerAotContribution() {
this.applicationContext = new AnnotationConfigApplicationContext();
this.applicationContext.register(ImportNotAccessibleContainer.class);
compile((freshContext, compiled) -> {
PostgreSQLContainer<?> container = freshContext.getBean(PostgreSQLContainer.class);
assertThat(container).isSameAs(ImportNotAccessibleContainer.container);
assertThat(freshContext.getEnvironment().getProperty("container.port", Integer.class))
.isEqualTo(ImportNotAccessibleContainer.container.getFirstMappedPort());
});
}

@SuppressWarnings("unchecked")
private void compile(BiConsumer<GenericApplicationContext, Compiled> result) {
ClassName className = processAheadOfTime();
TestCompiler.forSystem().with(this.generationContext).compile((compiled) -> {
try (GenericApplicationContext context = new GenericApplicationContext()) {
new TestcontainersLifecycleApplicationContextInitializer().initialize(context);
ApplicationContextInitializer<GenericApplicationContext> initializer = compiled
.getInstance(ApplicationContextInitializer.class, className.toString());
initializer.initialize(context);
context.refresh();
result.accept(context, compiled);
}
});
}

private ClassName processAheadOfTime() {
ClassName className = new ApplicationContextAotGenerator().processAheadOfTime(this.applicationContext,
this.generationContext);
this.generationContext.writeGeneratedContent();
return className;
}

@ImportTestcontainers
static class ImportWithoutValue {

Expand Down Expand Up @@ -196,4 +289,31 @@ void containerProperties() {

}

@ImportTestcontainers
static class CustomPostgreSQLContainerDefinitions {

static CustomPostgreSQLContainer container = new CustomPostgreSQLContainer();

}

static class CustomPostgreSQLContainer extends PostgreSQLContainer<CustomPostgreSQLContainer> {

CustomPostgreSQLContainer() {
super("postgres:14");
}

}

@ImportTestcontainers
static class ImportNotAccessibleContainer {

private static final PostgreSQLContainer<?> container = TestImage.container(PostgreSQLContainer.class);

@DynamicPropertySource
private static void containerProperties(DynamicPropertyRegistry registry) {
registry.add("container.port", container::getFirstMappedPort);
}

}

}
145 changes: 145 additions & 0 deletions ...org/springframework/boot/testcontainers/context/DynamicPropertySourceMethodsImporter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,31 @@

import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.util.Map;
import java.util.Set;

import org.springframework.aot.generate.AccessControl;
import org.springframework.aot.generate.GeneratedClass;
import org.springframework.aot.generate.GeneratedMethod;
import org.springframework.aot.generate.GenerationContext;
import org.springframework.aot.generate.MethodReference.ArgumentCodeGenerator;
import org.springframework.aot.hint.ExecutableMode;
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution;
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor;
import org.springframework.beans.factory.aot.BeanFactoryInitializationCode;
import org.springframework.beans.factory.aot.BeanRegistrationExcludeFilter;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.beans.factory.support.RegisteredBean;
import org.springframework.beans.factory.support.RootBeanDefinition;
import org.springframework.boot.testcontainers.properties.TestcontainersPropertySource;
import org.springframework.core.MethodIntrospector;
import org.springframework.core.annotation.MergedAnnotations;
import org.springframework.core.env.ConfigurableEnvironment;
import org.springframework.core.env.Environment;
import org.springframework.javapoet.CodeBlock;
import org.springframework.test.context.DynamicPropertyRegistry;
import org.springframework.test.context.DynamicPropertySource;
import org.springframework.util.Assert;
Expand Down Expand Up @@ -56,6 +74,16 @@ void registerDynamicPropertySources(BeanDefinitionRegistry beanDefinitionRegistr
ReflectionUtils.makeAccessible(method);
ReflectionUtils.invokeMethod(method, null, dynamicPropertyRegistry);
});

String beanName = "importTestContainer.%s.%s".formatted(DynamicPropertySource.class.getName(), definitionClass);
if (!beanDefinitionRegistry.containsBeanDefinition(beanName)) {
RootBeanDefinition bd = new RootBeanDefinition(DynamicPropertySourceMetadata.class);
bd.setInstanceSupplier(() -> new DynamicPropertySourceMetadata(definitionClass, methods));
bd.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
bd.setAutowireCandidate(false);
bd.setAttribute(DynamicPropertySourceMetadata.class.getName(), true);
beanDefinitionRegistry.registerBeanDefinition(beanName, bd);
}
}

private boolean isAnnotated(Method method) {
Expand All @@ -71,4 +99,121 @@ private void assertValid(Method method) {
+ "' must accept a single DynamicPropertyRegistry argument");
}

private record DynamicPropertySourceMetadata(Class<?> definitionClass, Set<Method> methods) {
}

static class DynamicPropertySourceMetadataBeanRegistrationExcludeFilter implements BeanRegistrationExcludeFilter {

@Override
public boolean isExcludedFromAotProcessing(RegisteredBean registeredBean) {
return registeredBean.getMergedBeanDefinition().hasAttribute(DynamicPropertySourceMetadata.class.getName());
}

}

/**
* {@link BeanFactoryInitializationAotProcessor} that generates methods for each
* annotated {@link DynamicPropertySource} method.
*/
static class DynamicPropertySourceBeanFactoryInitializationAotProcessor
implements BeanFactoryInitializationAotProcessor {

@Override
public BeanFactoryInitializationAotContribution processAheadOfTime(
ConfigurableListableBeanFactory beanFactory) {
Map<String, DynamicPropertySourceMetadata> metadata = beanFactory
.getBeansOfType(DynamicPropertySourceMetadata.class);
if (metadata.isEmpty()) {
return null;
}
return new AotContibution(metadata);
}

private static final class AotContibution implements BeanFactoryInitializationAotContribution {

private final Map<String, DynamicPropertySourceMetadata> metadata;

private AotContibution(Map<String, DynamicPropertySourceMetadata> metadata) {
this.metadata = metadata;
}

@Override
public void applyTo(GenerationContext generationContext,
BeanFactoryInitializationCode beanFactoryInitializationCode) {
GeneratedMethod initializerMethod = beanFactoryInitializationCode.getMethods()
.add("registerDynamicPropertySources", (code) -> {
code.addJavadoc("Registers {@code @DynamicPropertySource} properties");
code.addParameter(ConfigurableEnvironment.class, "environment");
code.addParameter(DefaultListableBeanFactory.class, "beanFactory");
code.addModifiers(javax.lang.model.element.Modifier.PRIVATE,
javax.lang.model.element.Modifier.STATIC);
code.addStatement("$T dynamicPropertyRegistry = $T.attach(environment, beanFactory)",
DynamicPropertyRegistry.class, TestcontainersPropertySource.class);
this.metadata.forEach((name, metadata) -> {
GeneratedMethod dynamicPropertySourceMethod = generateMethods(generationContext, metadata);
code.addStatement(dynamicPropertySourceMethod.toMethodReference()
.toInvokeCodeBlock(ArgumentCodeGenerator.of(DynamicPropertyRegistry.class,
"dynamicPropertyRegistry")));
});
});
beanFactoryInitializationCode.addInitializer(initializerMethod.toMethodReference());
}

// Generates a new class in definition class package and invokes
// all @DynamicPropertySource methods.
private GeneratedMethod generateMethods(GenerationContext generationContext,
DynamicPropertySourceMetadata metadata) {
Class<?> definitionClass = metadata.definitionClass();
GeneratedClass generatedClass = generationContext.getGeneratedClasses()
.addForFeatureComponent(DynamicPropertySource.class.getSimpleName(), definitionClass,
(code) -> code.addModifiers(javax.lang.model.element.Modifier.PUBLIC));
return generatedClass.getMethods().add("registerDynamicPropertySource", (code) -> {
code.addJavadoc("Registers {@code @DynamicPropertySource} properties for class '$L'",
definitionClass.getName());
code.addParameter(DynamicPropertyRegistry.class, "dynamicPropertyRegistry");
code.addModifiers(javax.lang.model.element.Modifier.PUBLIC,
javax.lang.model.element.Modifier.STATIC);
metadata.methods().forEach((method) -> {
GeneratedMethod generateMethod = generateMethod(generationContext, generatedClass,
definitionClass, method);
code.addStatement(generateMethod.toMethodReference()
.toInvokeCodeBlock(ArgumentCodeGenerator.of(DynamicPropertyRegistry.class,
"dynamicPropertyRegistry")));
});
});
}

// If the method is inaccessible, the reflection will be used; otherwise,
// direct call to the method will be used.
private static GeneratedMethod generateMethod(GenerationContext generationContext,
GeneratedClass generatedClass, Class<?> definitionClass, Method method) {
return generatedClass.getMethods().add(method.getName(), (code) -> {
code.addJavadoc("Register {@code @DynamicPropertySource} for method '$L.$L'",
method.getDeclaringClass().getName(), method.getName());
code.addModifiers(javax.lang.model.element.Modifier.PRIVATE,
javax.lang.model.element.Modifier.STATIC);
code.addParameter(DynamicPropertyRegistry.class, "dynamicPropertyRegistry");
if (AccessControl.forMember(method).isAccessibleFrom(generatedClass.getName())) {
code.addStatement(
CodeBlock.of("$T.$L(dynamicPropertyRegistry)", definitionClass, method.getName()));
}
else {
generationContext.getRuntimeHints().reflection().registerMethod(method, ExecutableMode.INVOKE);
code.addStatement("$T method = $T.findMethod($T.class, $S, $T.class)", Method.class,
ReflectionUtils.class, definitionClass, method.getName(),
DynamicPropertyRegistry.class);
code.addStatement("$T.notNull(method, $S)", Assert.class,
"Method '" + method.getName() + "' is not found");
code.addStatement("$T.makeAccessible(method)", ReflectionUtils.class);
code.addStatement("$T.invokeMethod(method, null, dynamicPropertyRegistry)",
ReflectionUtils.class);
}
});

}

}

}

}
Loading

0 comments on commit 91e3273

Please sign in to comment.