postgres junit extension

Integration tests are useful for identifying implementation caveats in real infrastructure components. One such component that I often use is a postgresql database. To write these integration tests, I use testcontainers to create a lightweight instance of postgresql.

The tutorials suggested by testcontainers are good, but I still need to create some boilerplate code to configure the database connections.

Since I primarily use spring-based applications, I could use the Spring Boot Testcontainers integration. However, this would require me to copy and paste the configuration for each of my integration tests, which can be cumbersome.

To avoid writing excessive boilerplate code, I decided to create a custom junit extension.

Moreover, I’m also using liquibase to perform database migrations, so I also want this extension to support this case.

Maven configuration

Add those dependencies:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-testcontainers</artifactId>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.testcontainers</groupId>
    <artifactId>postgresql</artifactId>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.postgresql</groupId>
    <artifactId>postgresql</artifactId>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-jpa</artifactId>
    <scope>test</scope>
</dependency>

Spring configuration

First, I would need to create a PostgresTestcontainersConfig, a spring test configuration file:

import org.springframework.boot.test.context.TestConfiguration;
import org.springframework.boot.testcontainers.service.connection.ServiceConnection;
import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.test.annotation.DirtiesContext;
import org.springframework.test.context.DynamicPropertyRegistrar;
import org.testcontainers.containers.PostgreSQLContainer;
 
import lombok.extern.slf4j.Slf4j;
 
@Slf4j
@TestConfiguration(proxyBeanMethods = false)
class PostgresTestcontainersConfig {
 
    /**
     * The {@link ServiceConnection} will make Spring Boot's auto-configuration
     * consume the details of a service connection and use them to establish a
     * connection to the PostgreSQL container.
     *
     * @see <a href="https://docs.spring.io/spring-boot/reference/testing/testcontainers.html#testing.testcontainers.service-connections">Spring Boot Testcontainers integration</a>
     */
    @Bean
    @ServiceConnection
    PostgreSQLContainer<?> postgresContainer() {
        return PostgresExtension.postgreSQLContainer;
    }
 
    /**
    * Dynamically configures Spring properties to enable or disable database
    * migration based on the fields of the
    * {@link RunPostgresContainer @RunPostgresContainer} annotation. This
    * includes setting Spring application properties.
    * <br>
    * This configuration ensures that the appropriate properties are set to
    * control the behavior of database migrations during tests.
    * <br>
    * If you discover that tests in subclasses fail because the dynamic
    * properties change between subclasses, you may need to annotate your base
    * class with {@link DirtiesContext @DirtiesContext} to ensure that each
    * subclass gets its own {@link ApplicationContext} with the correct dynamic
    * properties.
    *
    * @see <a href="https://contribute.liquibase.com/extensions-integrations/directory/integration-docs/springboot/configuration/">Liquibase Spring Boot integration</a>
    * @see <a href="https://docs.spring.io/spring-framework/reference/testing/testcontext-framework/ctx-management/dynamic-property-sources.html">Context configuration with Dynamic Property Sources</a>
    */
    @Bean
    DynamicPropertyRegistrar databaseMigrationRegistrar() {
        return registry -> {
            var postgreSQLContainerProperties = PostgresExtension.postgreSQLContainerProperties;
            if (postgreSQLContainerProperties == null) {
                log.warn("The PostgreSQL container should have been initialized before spring application!");
                return;
            }
 
            if (postgreSQLContainerProperties.runDatabaseMigration()) {
                log.info("Database migration ENABLED.");
                registry.add("spring.liquibase.enabled", () -> "true");
 
                String liquibaseChangeLogPath = postgreSQLContainerProperties.liquibaseChangeLogPath();
                if (!liquibaseChangeLogPath.isBlank()) {
                    log.info("Using Liquibase change log: {}.", liquibaseChangeLogPath);
                    registry.add("spring.liquibase.change-log", () -> liquibaseChangeLogPath);
                }
            } else {
                log.info("Database migration DISABLED.");
                registry.add("spring.liquibase.enabled", () -> "false");
            }
        };
    }
}

• The first bean postgresContainer is used to have spring do its magic to connect the testcontainers instance with the Datasource, and set the right spring properties (e.g. spring.datasource.url, spring.datasource.driver-class-name).
• The second bean is used set the spring properties for performing the database migration if the consumer of the extension wants it.

Junit extension

Then, the junit extension PostgresExtension:

import static java.util.function.Predicate.not;
 
import java.lang.annotation.Annotation;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.extension.BeforeAllCallback;
import org.junit.jupiter.api.extension.ExtensionContext;
import org.junit.jupiter.api.extension.ParameterContext;
import org.junit.jupiter.api.extension.ParameterResolutionException;
import org.junit.jupiter.api.extension.ParameterResolver;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.utility.DockerImageName;
 
/**
 * DO NOT USE this extension directly.
 * <br>
 * Please use {@link RunPostgresContainer @RunPostgresContainer} instead.
 */
@Slf4j
final class PostgresExtension implements BeforeAllCallback, ParameterResolver {
 
    static final String DEFAULT_DOCKER_IMAGE_NAME = "postgres:latest";
 
    static PostgreSQLContainer<?> postgreSQLContainer;
    static PostgreSQLContainerProperties postgreSQLContainerProperties;
 
    @Override
    public void beforeAll(ExtensionContext context) throws Exception {
        var properties = findPostgreSQLContainerProperties(context);
        // (Re-)create the PostgreSQL container if the annotation has different
        // values from a previous test.
        if (postgreSQLContainerProperties == null || !postgreSQLContainerProperties.equals(properties)) {
            postgreSQLContainerProperties = properties;
            log.info(
                    "Starting new PostgreSQL container with the following properties:\n{}",
                    postgreSQLContainerProperties);
 
            // Testcontainers takes care of removing any created resources
            // (containers, volumes, networks etc.) automatically after the test
            // execution is complete by using the Ryuk sidecar container.
            // See https://testcontainers.com/getting-started/#benefits-of-using-testcontainers.
            postgreSQLContainer = new PostgreSQLContainer<>(postgreSQLContainerProperties.dockerImageName());
            postgreSQLContainer = customizePostgreSQLContainer(postgreSQLContainerProperties, postgreSQLContainer);
            postgreSQLContainer.start();
        } else {
            log.info("Using existing PostgreSQL container.");
        }
    }
 
    @Override
    public boolean supportsParameter(ParameterContext parameterContext, ExtensionContext extensionContext)
            throws ParameterResolutionException {
        return parameterContext.getParameter().getType() == PostgreSQLContainer.class;
    }
 
    @Override
    public Object resolveParameter(ParameterContext parameterContext, ExtensionContext extensionContext)
            throws ParameterResolutionException {
        return postgreSQLContainer;
    }
 
    private PostgreSQLContainerProperties findPostgreSQLContainerProperties(ExtensionContext context) {
        return findAnnotationRecursively(context.getRequiredTestClass(), RunPostgresContainer.class)
                .map(PostgreSQLContainerProperties::from)
                .orElseThrow(() -> new IllegalArgumentException(
                        "Do not use PostgresExtension directly, use @RunPostgresContainer instead!"));
    }
 
    private <A extends Annotation> Optional<A> findAnnotationRecursively(Class<?> clazz, Class<A> annotationClass) {
        if (clazz.isAnnotationPresent(annotationClass)) {
            return Optional.of(clazz.getAnnotation(annotationClass));
        }
        Class<?> enclosingClass = clazz.getEnclosingClass();
        if (enclosingClass != null) {
            return findAnnotationRecursively(enclosingClass, annotationClass);
        }
        return Optional.empty();
    }
 
    private PostgreSQLContainer<?> customizePostgreSQLContainer(
            PostgreSQLContainerProperties postgreSQLContainerProperties, PostgreSQLContainer<?> postgreSQLContainer) {
        if (!postgreSQLContainerProperties.containerInitScripts().isEmpty()) {
            return postgreSQLContainer.withInitScripts(postgreSQLContainerProperties.containerInitScripts());
        }
        return postgreSQLContainer;
    }
 
    record PostgreSQLContainerProperties(
            String fullImageName,
            List<String> containerInitScripts,
            boolean runDatabaseMigration,
            String liquibaseChangeLogPath) {
        static PostgreSQLContainerProperties from(RunPostgresContainer runPostgresContainer) {
            return new PostgreSQLContainerProperties(
                    runPostgresContainer.dockerImageName(),
                    Arrays.asList(runPostgresContainer.containerInitScripts()).stream()
                            .filter(not(String::isBlank))
                            .toList(),
                    runPostgresContainer.runDatabaseMigration(),
                    runPostgresContainer.liquibaseChangeLogPath());
        }
 
        private DockerImageName dockerImageName() {
            if (fullImageName.isBlank()) {
                return DockerImageName.parse(DEFAULT_DOCKER_IMAGE_NAME);
            }
            return DockerImageName.parse(fullImageName);
        }
    }
}

I use a PostgreSQLContainerProperties to store the properties of the extension and adapt its behavior accordingly. For example, I can add some postgres init scripts to execute on container startup, perform database migrations with liquibase, and more.

So instead of checking the nullity of the postgresContainer instance, I use this PostgreSQLContainerProperties to check whether the container should be re-created or not. This way, as a consumer of the extension, I would not be confused on the test behavior. So that still means that if I want to have quick tests, I should use the same properties for all of my integration tests.

Dedicated annotation for bootstraping

Finally, an annotation RunPostgresContainer to bootstrap those two files:

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import org.junit.jupiter.api.extension.ExtendWith;
import org.springframework.boot.test.autoconfigure.jdbc.AutoConfigureTestDatabase;
import org.springframework.boot.test.autoconfigure.jdbc.AutoConfigureTestDatabase.Replace;
import org.springframework.boot.test.autoconfigure.orm.jpa.DataJpaTest;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Import;
import org.springframework.test.annotation.DirtiesContext;
import org.springframework.test.context.cache.DefaultContextCache;
import org.testcontainers.Testcontainers;
import org.testcontainers.containers.Container;
import org.testcontainers.containers.PostgreSQLContainer;
 
/**
 * Annotation used to start a {@link PostgreSQLContainer} and perform database
 * migration for your tests.
 * <br>
 * This annotation is designed for use with Spring-based tests, such as
 * slice tests using annotations like {@link DataJpaTest @DataJpaTest}, or full
 * Spring applications using {@link SpringBootTest @SpringBootTest}.
 *
 * <h2>Usage Example:</h2>
 *
 * <pre>
 * &#64;RunPostgresContainer(
 *      dockerImageName = "postgres:16-alpine",
 *      containerInitScripts = {"db/init_1.sql", "db/init_2.sql"},
 *      runDatabaseMigration = true,
 *      liquibaseChangeLogPath = "classpath:db/custom/db.changelog-custom.xml"
 * )
 * // You can use slice tests with &#64;DataJpaTest, &#64;DataJdbcTest, ...
 * // or directly &#64;SpringBootTest.
 * &#64;DataJpaTest
 * class AnIntegrationTest {
 *     &#64;Autowired
 *     private MyRepository myRepository;
 *
 *     &#64;BeforeEach
 *     void beforeEach() {
 *         // Clean up your data if you do not rollback each data test.
 *         myRepository.deleteAll();
 *     }
 *
 *     &#64;Test
 *     void testDatabaseInteraction(
 *         // You can inject the container as a method parameter if you need to
 *         // get some information from it.
 *         PostgreSQLContainer<?> postgreSQLContainer
 *     ) {
 *         // Your test code here
 *     }
 * }
 * </pre>
 *
 * /!\ This annotation is designed to start a single container for several test
 * classes to speed up test executions. Therefore, be aware that your tests
 * will share a {@link PostgreSQLContainer}!
 * <br>
 * However, if the annotation has different field values, a new
 * {@link PostgreSQLContainer} will be created.
 * Do not forget to add {@link DirtiesContext @DirtiesContext} if you are using
 * {@link SpringBootTest @SpringBootTest} to ensure Spring test will not use its
 * cache ({@link DefaultContextCache}) and re-use the same spring context.
 * Otherwise, your database connection will not be updated with the new
 * {@link PostgreSQLContainer}.
 * <br>
 * This is not thread-safe, so be careful when injecting
 * {@link PostgreSQLContainer} in your test method parameter and running your
 * tests in parallel.
 * <br>
 * If you want to create a custom {@link PostgreSQLContainer} with specific
 * configurations for your tests, it's best to use {@link Testcontainers} and
 * {@link Container @Container} from Testcontainers directly, as suggested by the
 * official Testcontainers documentation.
 *
 * @see <a href="https://java.testcontainers.org/test_framework_integration/junit_5/">Testcontainers JUnit 5 Integration</a>
 * @see <a href="https://java.testcontainers.org/test_framework_integration/manual_lifecycle_control/#singleton-containers">Singleton Containers</a>
 * @see <a href="https://docs.spring.io/spring-framework/reference/testing/testcontext-framework/ctx-management/caching.html">Spring context caching</a>
 */
@Target({ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Import(PostgresTestcontainersConfig.class)
@ExtendWith(PostgresExtension.class)
// Ensure spring test won't use an in-memory db.
@AutoConfigureTestDatabase(replace = Replace.NONE)
public @interface RunPostgresContainer {
 
    /**
     * @return the Docker image name to use to start {@link PostgreSQLContainer}
     */
    String dockerImageName() default PostgresExtension.DEFAULT_DOCKER_IMAGE_NAME;
 
    /**
     * You can use this field to add PostgreSQL extensions, create custom schemas, etc...
     *
     * @return ordered collection of SQL scripts for {@link PostgreSQLContainer} initialization
     */
    String[] containerInitScripts() default "";
 
    /**
     * @return {@code true} if you want to perform database migration on the {@link PostgreSQLContainer}.
     */
    boolean runDatabaseMigration() default false;
 
    /**
     * If not set, will use the default value of {@code spring.liquibase.change-log}, which is
     * {@code classpath:/db/changelog/db.changelog-master.yaml}.
     * <br>
     * This field is used only if {@link #runDatabaseMigration()} is set to {@code true}.
     *
     * @return the path to the Liquibase change logs path to use for the
     *         database migration
     * @see https://docs.spring.io/spring-boot/appendix/application-properties/index.html#application-properties.data-migration.spring.liquibase.change-log
     */
    String liquibaseChangeLogPath() default "";
}

Usage

@SpringBootTest
@RunPostgresContainer(
  dockerImageName = "postgres:16-alpine",
  containerInitScripts = "db/init.sql",
  runDatabaseMigration = true,
  liquibaseChangeLogPath = "db/db-changelog.xml"
)
class MyAwesomeIT {
    @Autowired
    private MyAwesomeRepository myAwesomeRepository;
 
    @BeforeEach
    void beforeEach() {
        // Clean up your data if you do not rollback each data test.
        myAwesomeRepository.deleteAll();
    }
 
     @Test
     void testDatabaseInteraction(
         // You can inject the container as a method parameter if you need to
         // get some information from it.
         PostgreSQLContainer<?> postgreSQLContainer
     ) {
         // Your test code here
     }
}