Integration Testing AI Prompts With Ollama and Spring TestContainers

AI features are becoming common in modern applications. If your Spring Boot app uses large language models (LLMs), it’s important to test how those models respond to real prompts. This helps you catch issues early and keeps your app reliable.

In this article, you’ll learn how to write integration tests for AI prompts using Spring TestContainers and Ollama. You’ll see how to set up your environment, write prompt tests, and apply good testing practices - all using standard JUnit and Spring Boot.

Why Test AI Prompts?

Prompt responses from LLMs are not always predictable. Even small changes in the model or configuration can cause different outputs. Testing AI prompts helps you:

• Check if responses are valid
• Catch regressions when models change
• Ensure APIs using LLMs work as expected

What Is Spring TestContainers?

Spring TestContainers is a library that makes it easier to use Testcontainers in Spring applications. It uses annotations to manage Docker containers during tests. You don’t need to write any container setup code.

With built-in support for Ollama, Spring TestContainers can automatically:

• Start a Docker container with the Ollama runtime
• Pull the right model version
• Inject a working ChatClient into your Spring tests
• Clean up the container after the test finishes

Getting Started

First, add the necessary dependencies to your project:

Then, use the @EnableOllamaContainer annotation on your test class:

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

This sets up the Ollama container and downloads the AI model.

Writing Your First Prompt Test

Here’s an example that tests a simple prompt:

@SpringBootTest
@EnableOllamaContainer(model = "llama3:latest")
public class BasicOllamaTest {

    @Autowired
    private ChatClient.Builder chatClientBuilder;

    private ChatClient chatClient;

    @BeforeEach
    void init() {
        this.chatClient = chatClientBuilder.build();
    }

    @Test
    void testCapitalPrompt() {
        String response = chatClient.prompt()
            .user("What is the capital of France?")
            .call()
            .content();
        assertTrue(response.toLowerCase().contains("paris"));
    }
}

Testing for Deterministic Output

You can lower the temperature setting to get more consistent results:

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

This is helpful for math or logic prompts:

@Test
void testMathPrompt() {
    String response = chatClient.prompt()
        .user("What is 1 + 2? Give the value only.")
        .call()
        .content();
    assertTrue(response.contains("3"));
}

Testing Multiple Prompts

You can test several prompts in one class using @ParameterizedTest:

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

Testing REST APIs That Use Ollama

If your controller calls an Ollama-based service, you can test the endpoint too:

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
@EnableOllamaContainer(model = "llama3:latest")
public class ChatApiTest {
    @LocalServerPort
    private int port;

    @Autowired
    private TestRestTemplate restTemplate;

    @Test
    void testChatEndpoint() {
        String url = "http://localhost:" + port + "/api/chat?message=What is 1+1?";
        String response = restTemplate.getForObject(url, String.class);
        assertTrue(response.contains("2"));
    }
}

How It Works

Here’s what happens behind the scenes when you use @EnableOllamaContainer:

1. Spring detects the annotation and starts an Ollama container
2. It pulls the required model version
3. Spring wires in the connection details for you
4. Your test runs against the live Ollama container
5. The container stops and is removed after the test

This process uses the standard Testcontainers Java library but wraps it in a simpler API tailored for Spring Boot.

Benefits of Using Spring TestContainers With Ollama

Spring TestContainers handles much of the setup for you. Here are the key benefits:

• Clean test isolation – Each test uses a fresh Ollama container
• No manual setup – Just use an annotation to configure the container
• Automatic model management – The right model is pulled and applied
• Model caching – Spring TestContainers mounts a shared volume so models persist across tests
• Works with JUnit 5 – Integrates smoothly into your test lifecycle

Best Practices for Prompt Testing

Here are a few ways to make your tests more stable and maintainable:

• Use fixed model versions – Avoid using latest tags in tests. Instead, use something like llama3:8b.
• Set a low temperature for consistent output – Use temperature = 0.1 if you expect the same response each time.
• Test for important phrases, not exact matches – Match only the part of the response that matters.
• Use parameterized tests for coverage – You can run multiple prompts in one class with @CsvSource.

Supported Models

Ollama supports a variety of open LLMs. You can use any of these in your tests:

• llama3:latest – Meta’s Llama 3
• llama2:latest – Meta’s Llama 2
• mistral:latest – Mistral
• phi3:latest – Microsoft’s Phi
• gemma:latest – Google’s Gemma

More are available in the Ollama model library.

Example Project

For a full working example, check out the springboot-ollama example in the repository.

Summary

Testing AI prompt logic is essential. Spring TestContainers and Ollama let you run real prompt tests with minimal setup. You write clean, reliable tests that work locally and in CI.

Head to the Spring TestContainers GitHub repo to get started.