AI Prompt Testing with Ollama Spring TestContainers

This guide shows you how to use Spring TestContainers to write integration tests for AI prompts with Ollama. It includes step-by-step instructions, examples, and tips for setting up and running tests in your Spring Boot app.

Introduction

Testing AI prompt interactions helps ensure your app’s AI features work as expected. Spring TestContainers makes this easier by letting you spin up Ollama containers with a simple annotation.

Why use Spring TestContainers for Ollama testing:

Easy setup: Add @EnableOllamaContainer to your test class.
No manual config: Spring auto-configures the container connection.
Test isolation: Each test gets a clean Ollama instance.
Model handling: The right AI model is pulled and managed for you.
JUnit 5 ready: Works smoothly with the JUnit 5 test lifecycle.
Faster runs: Containers and models are reused and cached for repeat tests.

Getting Started

1. Add Dependencies

First, add the necessary dependencies to your project:

Maven


<dependencies>
    <!-- Spring AI -->
    <dependency>
        <groupId>org.springframework.ai</groupId>
        <artifactId>spring-ai-ollama-spring-boot-starter</artifactId>
        <version>0.8.0</version>
    </dependency>
 
    <!-- Spring TestContainers -->
    <dependency>
        <groupId>io.flowinquiry</groupId>
        <artifactId>spring-testcontainers-ollama</artifactId>
        <version><!-- Replace with the latest version --></version>
        <scope>test</scope>
    </dependency>
 
    <!-- JUnit and Spring Test -->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-test</artifactId>
        <scope>test</scope>
    </dependency>
</dependencies>

🔍 Find the latest version: Search on Maven Central

Writing AI Prompt Tests with Ollama

Basic Ollama Test

To create a basic integration test with Ollama:


import io.flowinquiry.testcontainers.ai.EnableOllamaContainer;
import io.flowinquiry.testcontainers.ai.OllamaOptions;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
 
@SpringBootTest
@EnableOllamaContainer(
    dockerImage = "ollama/ollama",
    version = "0.9.0",
    model = "llama3:latest",
    options = @OllamaOptions(temperature = "0.7", topP = "0.5"))
public class BasicOllamaTest {
 
    @Autowired
    private ChatClient.Builder chatClientBuilder;
 
    private ChatClient chatClient;
 
    @BeforeEach
    public void init() {
        this.chatClient = this.chatClientBuilder.build();
    }
 
    @Test
    public void testSimplePrompt() {
        String prompt = "What is the capital of France?";
        String response = chatClient.prompt().user(prompt).call().content();
 
        assertNotNull(response);
        assertTrue(response.contains("Paris"));
    }
}

Testing Deterministic Responses

For tests that require more deterministic responses, you can use a lower temperature setting:


@SpringBootTest
@EnableOllamaContainer(
    model = "llama3:latest",
    options = @OllamaOptions(temperature = "0.1", topP = "0.5"))
public class DeterministicResponseTest {
 
    @Autowired
    private ChatClient.Builder chatClientBuilder;
 
    private ChatClient chatClient;
 
    @BeforeEach
    public void init() {
        this.chatClient = this.chatClientBuilder.build();
    }
 
    @Test
    public void testMathQuestion() {
        String prompt = "What is the result of 1+2? Give the value only";
        String response = chatClient.prompt().user(prompt).call().content();
 
        assertTrue(response.contains("3"));
    }
}

Parameterized Tests

You can use JUnit’s parameterized tests to test multiple prompts:


@SpringBootTest
@EnableOllamaContainer(model = "llama3:latest")
public class ParameterizedOllamaTest {
 
    @Autowired
    private ChatClient.Builder chatClientBuilder;
 
    private ChatClient chatClient;
 
    @BeforeEach
    public void init() {
        this.chatClient = this.chatClientBuilder.build();
    }
 
    @ParameterizedTest
    @CsvSource({
        "What is the result of 1+2? Give the value only, 3",
        "How many letter 'r' in the word 'Hello'? Give the value only, 0"
    })
    public void testPrompts(String prompt, String expectedResult) {
        String content = chatClient.prompt().user(prompt).call().content();
 
        assertNotNull(content);
        assertFalse(content.isEmpty());
        assertTrue(content.contains(expectedResult));
    }
}

Testing REST API Endpoints

If your application exposes REST endpoints that use Ollama, you can test them with TestRestTemplate:


@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
@EnableOllamaContainer(model = "llama3:latest")
public class ChatControllerTest {
 
    @LocalServerPort
    private int port;
 
    @Autowired
    private TestRestTemplate restTemplate;
 
    @Test
    public void testChatEndpoint() {
        String url = "http://localhost:" + port + "/api/chat?message=What is 1+1?";
        String response = restTemplate.getForObject(url, String.class);
 
        assertNotNull(response);
        assertTrue(response.contains("2"));
    }
}

How It Works

Spring TestContainers for Ollama works by:

Container Creation: When your test starts, Spring TestContainers detects the @EnableOllamaContainer annotation and creates an Ollama container with the specified Docker image and version.
Container Startup: The container is started before your test runs.
Model Pulling: The specified AI model is automatically pulled if it’s not already available in the container.
Spring Configuration: Spring’s environment is automatically configured with the container’s connection details and model options.
Test Execution: Your test runs against the containerized Ollama instance.
Container Cleanup: After your test completes, the container is automatically stopped and removed.

Behind the scenes, Spring TestContainers uses the Testcontainers library to manage the Docker containers, but provides a much simpler API through annotations.

Advanced Configuration

Custom Docker Images

You can specify a custom Docker image by using the dockerImage and version parameters:

The annotation would look like this:

@EnableOllamaContainer(dockerImage = "custom/ollama", version = "1.0.0", model = "llama3:latest")

Model Options

You can configure the AI model’s behavior using the options parameter:

The annotation would look like this:

@EnableOllamaContainer(model = "llama3:latest", options = @OllamaOptions(temperature = "0.7", topP = "0.9"))

Where:

temperature: Controls randomness (0.0 to 1.0)
topP: Controls diversity (0.0 to 1.0)

Available Models

Ollama supports various models. Some commonly used ones include:

llama3:latest - Meta’s Llama 3 model
llama2:latest - Meta’s Llama 2 model
mistral:latest - Mistral AI’s model
phi3:latest - Microsoft’s Phi-3 model
gemma:latest - Google’s Gemma model

Check the Ollama model library for more available models.

Best Practices

1. Use Specific Model Versions

Always specify the exact model version to ensure test reproducibility:

For example:

@EnableOllamaContainer(model = "llama3:8b")

2. Set Appropriate Temperature

For deterministic tests, use a low temperature value:

For example:

@EnableOllamaContainer(model = "llama3:latest", options = @OllamaOptions(temperature = "0.1"))

3. Test for Content Patterns

Instead of expecting exact responses, test for patterns or key information:


@Test
public void testCapitalQuestion() {
    String response = chatClient.prompt().user("What is the capital of France?").call().content();
 
    // Don't test for exact match
    // Bad: assertEquals("The capital of France is Paris.", response);
 
    // Instead, test for key information
    assertTrue(response.toLowerCase().contains("paris"));
}

4. Use Parameterized Tests

Use parameterized tests to test multiple prompts with expected results:


@ParameterizedTest
@CsvSource({
    "What is 2+2?, 4",
    "What is the capital of Japan?, Tokyo",
    "Who wrote Romeo and Juliet?, Shakespeare"
})
public void testFactualQuestions(String prompt, String expectedContent) {
    String response = chatClient.prompt().user(prompt).call().content();
    assertTrue(response.contains(expectedContent));
}

5. Cache Models Between Tests

To speed up test execution, consider caching models between test runs:

For example:

@EnableOllamaContainer(dockerImage = "ollama/ollama", version = "0.9.0", model = "llama3:latest")

The provider automatically binds /tmp/ollama-cache to /root/.ollama in the container, which caches downloaded models.

Example Application

For a complete example of using Ollama with Spring TestContainers, see the springboot-ollama example in the repository.

Conclusion

Spring TestContainers lets you test AI prompts with Ollama using a simple annotation. It runs containerized Ollama instances, so your tests use the same AI models as production. Each test stays isolated and repeatable.

This setup makes it easy to check how your app behaves when talking to large language models. You can write clear, reliable tests that cover real interactions.

AI Prompt Testing with Ollama Spring TestContainers

Introduction

Getting Started

1. Add Dependencies

Maven

Gradle

Writing AI Prompt Tests with Ollama

Basic Ollama Test

Testing Deterministic Responses

Parameterized Tests

Testing REST API Endpoints

How It Works

Advanced Configuration

Custom Docker Images

Model Options

Available Models

Best Practices

1. Use Specific Model Versions

2. Set Appropriate Temperature

3. Test for Content Patterns

4. Use Parameterized Tests

5. Cache Models Between Tests

Example Application

Conclusion