Writing a New Container Extension for Spring TestContainers

This guide explains how to create a new container extension for the Spring TestContainers library. It covers the core concepts, the extension process, and provides examples to help you implement your own container extensions.

Understanding SpringAwareContainerProvider

The SpringAwareContainerProvider is the abstract base class at the heart of the Spring TestContainers extension system. It provides the foundation for integrating TestContainers with Spring applications.

Purpose and Functionality

SpringAwareContainerProvider serves several key purposes:

Container Lifecycle Management: It implements the LifecycleAware interface to handle container startup and shutdown in sync with the Spring application context.
Spring Environment Integration: It provides mechanisms to apply container-specific configuration to the Spring environment.
Annotation-Based Configuration: It uses Java annotations to configure containers with parameters like Docker image and version.

Class Structure

The SpringAwareContainerProvider is a generic class with two type parameters:

A extends Annotation: The annotation type used to enable and configure the container
SELF extends GenericContainer<SELF>: The specific TestContainer type being managed


public abstract class SpringAwareContainerProvider<
        A extends Annotation, SELF extends GenericContainer<SELF>>
    implements LifecycleAware {
    // Implementation details
}

Key Methods

initContainerInstance(A enableContainerAnnotation): Initializes the container based on annotation parameters.
createContainer(): Abstract method that must be implemented by subclasses to create a specific container type.
start(): Starts the container when the Spring context is initialized.
stop(): Stops the container when the Spring context is closed.
getContainerType(): Abstract method that returns the type of container being managed.
applyTo(ConfigurableEnvironment environment): Abstract method that applies container-specific configuration to the Spring environment.

How to Extend SpringAwareContainerProvider

To create a new container extension, follow these steps:

1. Define an Annotation

Create an annotation that will be used to enable and configure your container:


@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Inherited
public @interface EnableMyContainer {
    String dockerImage() default "mycontainer/myimage";
    String version() default "latest";
 
    // Add any container-specific configuration options
    String myOption() default "defaultValue";
}

2. Create Your Container Provider

Extend the SpringAwareContainerProvider class to create your container provider:


public class MyContainerProvider
    extends SpringAwareContainerProvider<EnableMyContainer, MyContainer> {
 
    @Override
    public ContainerType getContainerType() {
        return MY_CONTAINER_TYPE; // Define this in ContainerType enum
    }
 
    @Override
    protected MyContainer createContainer() {
        return new MyContainer(dockerImage + ":" + version)
            .withExposedPorts(8080)
            .withEnv("MY_ENV_VAR", "value");
    }
 
    @Override
    public void applyTo(ConfigurableEnvironment environment) {
        Properties props = new Properties();
        props.put("my.container.url", "http://" + container.getHost() + ":" + container.getMappedPort(8080));
        props.put("my.container.option", enableContainerAnnotation.myOption());
 
        environment
            .getPropertySources()
            .addFirst(new PropertiesPropertySource("testcontainers", props));
    }
}

3. Register Your Provider

Create a service loader file to register your provider:

Create a file at src/main/resources/META-INF/services/io.flowinquiry.testcontainers.SpringAwareContainerProvider
Add the fully qualified name of your provider class to this file:


com.example.testcontainers.MyContainerProvider

4. Update ContainerType Enum (Optional)

If your container represents a new type, add it to the ContainerType enum:


public enum ContainerType {
    POSTGRESQL,
    MYSQL,
    OLLAMA,
    MY_CONTAINER_TYPE // Add your container type
}

Class Hierarchy and Relationships

Examples

Example 1: JDBC Container Extension

The SpringAwareJdbcContainerProvider is an example of extending SpringAwareContainerProvider for JDBC database containers:


public abstract class SpringAwareJdbcContainerProvider<T extends JdbcDatabaseContainer<T>>
    extends SpringAwareContainerProvider<EnableJdbcContainer, T> {
 
    @Override
    public final void applyTo(ConfigurableEnvironment environment) {
        Properties props = new Properties();
        props.put("spring.datasource.url", container.getJdbcUrl());
        props.put("spring.datasource.username", container.getUsername());
        props.put("spring.datasource.password", container.getPassword());
 
        environment
            .getPropertySources()
            .addFirst(new PropertiesPropertySource("testcontainers", props));
    }
}

Example 2: Ollama AI Container Extension

The OllamaContainerProvider is a concrete implementation for Ollama AI containers:


public class OllamaContainerProvider
    extends SpringAwareContainerProvider<EnableOllamaContainer, OllamaContainer> {
 
    @Override
    public ContainerType getContainerType() {
        return OLLAMA;
    }
 
    @Override
    protected OllamaContainer createContainer() {
        return new OllamaContainer(dockerImage + ":" + version)
            .withFileSystemBind("/tmp/ollama-cache", "/root/.ollama", READ_WRITE);
    }
 
    @Override
    public void start() {
        super.start();
        // Additional startup logic to pull AI models
        try {
            pullModelIfMissing(enableContainerAnnotation.model());
        } catch (IOException | InterruptedException e) {
            throw new RuntimeException(e);
        }
    }
 
    @Override
    public void applyTo(ConfigurableEnvironment environment) {
        Properties props = new Properties();
        props.put("spring.ai.ollama.base-url", container.getEndpoint());
        props.put("spring.ai.ollama.chat.model", enableContainerAnnotation.model());
        // Additional properties...
 
        environment
            .getPropertySources()
            .addFirst(new PropertiesPropertySource("testcontainers", props));
    }
}

Conclusion

By extending the SpringAwareContainerProvider class, you can create custom container extensions that integrate seamlessly with Spring applications. This approach allows you to leverage the power of TestContainers while maintaining the simplicity and elegance of Spring’s configuration system.