A Docker Compose Test Environment for LLM-Calling Code

#ai #machinelearning #tutorial #llm

The goal is not “a model in Docker”. It is that a developer who has just cloned the repository can run the full suite on a train, and that the suite gives the same answer there as it does in CI. That takes two inference services, not one.

What the environment has to provide

Three different kinds of test want three different things from the thing on the other end of the socket, and one service cannot be all three.

Deterministic responses on demand. Most tests need a specific response — a truncated JSON body, a 429 with a Retry-After header, an empty completion, a tool call with a malformed argument string. No model produces those on request. A stub does.
A real inference server. Some tests are about the wiring: does the SDK negotiate the stream correctly, does the request serialise, does the response deserialise, does the timeout fire. A stub that returns a hand-written body proves nothing about any of that because you wrote the body.
Nothing on the public internet. Both of the above must work offline and cost nothing, or they will be skipped locally and only run in CI, which is how a suite becomes something people wait on rather than something they use.

The compose file

Two services plus the application, one network, one named volume for the model weights so they survive a docker compose down.

# compose.test.yaml
services:
  llm:
    image: ollama/ollama:0.13.3
    volumes:
      - ollama-models:/root/.ollama
    healthcheck:
      test: ["CMD", "ollama", "list"]
      interval: 5s
      timeout: 3s
      retries: 20
    ports:
      - "11434:11434"

  llm-stub:
    build: ./test/stub
    environment:
      PORT: "8080"
    healthcheck:
      test: ["CMD", "wget", "-qO-", "http://localhost:8080/healthz"]
      interval: 2s
      timeout: 2s
      retries: 15
    ports:
      - "8080:8080"

  app:
    build: .
    depends_on:
      llm:
        condition: service_healthy
      llm-stub:
        condition: service_healthy
    environment:
      LLM_BASE_URL: "http://llm:11434/v1"
      LLM_STUB_URL: "http://llm-stub:8080/v1"
      LLM_API_KEY: "test-key-not-real"
      LLM_MODEL: "qwen2.5:0.5b"
    command: ["npm", "run", "test:integration"]

volumes:
  ollama-models:

Pin the image tag. latest means the environment changes under you between two runs of the same commit, which is the exact property a test environment exists to eliminate — and on this particular image it also changes which OpenAI-compatible fields are implemented, so a compatibility suite run against an unpinned tag is measuring an unknown.

Use condition: service_healthy rather than a bare depends_on. The plain form waits for the container to start, not for the process inside it to be listening, and the resulting failure is a connection refused on the first test of every cold run — which people work around by adding a sleep, which is slower and still flaky.

The stub service

Write it yourself; it is thirty lines and you need to control exactly what it returns. The one design decision worth making up front is how a test asks for a specific response. A header is the cleanest: the test sends x-stub-scenario: truncated-json and the stub switches on it, so a scenario is selected per request rather than by mutating shared server state that leaks between parallel tests.

// test/stub/server.ts
import { createServer } from "node:http";

const scenarios: Record<string, [number, unknown]> = {
  ok: [200, { choices: [{ index: 0, finish_reason: "stop",
        message: { role: "assistant", content: '{"items":[]}' } }] }],
  "truncated-json": [200, { choices: [{ index: 0, finish_reason: "length",
        message: { role: "assistant", content: '{"items":[{"name":"ap' } }] }],
  "empty": [200, { choices: [{ index: 0, finish_reason: "stop",
        message: { role: "assistant", content: "" } }] }],
  "rate-limited": [429, { error: { message: "slow down", type: "rate_limit" } }],
  "server-error": [503, { error: { message: "overloaded" } }],
};

createServer((req, res) => {
  if (req.url === "/healthz") { res.writeHead(200).end("ok"); return; }
  const key = String(req.headers["x-stub-scenario"] ?? "ok");
  const [status, body] = scenarios[key] ?? scenarios.ok;
  const headers: Record<string, string> = { "content-type": "application/json" };
  if (status === 429) headers["retry-after"] = "2";
  res.writeHead(status, headers).end(JSON.stringify(body));
}).listen(Number(process.env.PORT ?? 8080));

Give it a /healthz that does not depend on the scenario table, so compose’s health check cannot be broken by a change to the scenarios.

Resist the urge to make the stub clever. Every feature you add to it — a template engine, a recorded-response mode, a little routing DSL — is code that can be wrong, and a bug in the stub presents as a bug in the application. Its job is to return bytes you chose. If a test needs something the scenario table cannot express, that test probably wants an in-process interceptor rather than a network service, and only the tests that must cross a real socket belong here at all.

The scenarios worth having are the ones that are hard to produce any other way, and they are mostly failures rather than successes: the rate limit with a Retry-After, the 503, the response that never finishes, the stream that closes halfway through an event, the 200 carrying an error object in the body. A stub that only knows how to succeed is a stub that leaves your error paths untested, and error paths are where inference code spends its incidents.

The model pull is the hard part

Everything above works on the second run and hangs on the first, because the model is not downloaded yet and the first request pulls it — several gigabytes, inside whatever timeout your test framework has. Pull it explicitly, before the app service starts, and keep it in the named volume.

  llm-init:
    image: ollama/ollama:0.13.3
    depends_on:
      llm:
        condition: service_healthy
    environment:
      OLLAMA_HOST: "http://llm:11434"
    entrypoint: ["ollama", "pull", "qwen2.5:0.5b"]
    restart: "no"

Choose the smallest model that exercises the code path. A half-billion-parameter model gives poor answers and correct plumbing, and plumbing is all these tests are asserting on. Anything larger buys nothing here and costs every developer several gigabytes of disk and a much slower first run.

Two failure modes are worth knowing about before you meet them. The init container exits as soon as the pull finishes, which is correct but makes compose report a stopped container; restart: "no" keeps it from looping. And a pull that fails — a renamed tag, a registry outage — leaves the volume half-populated, so the next run appears to succeed and then serves a broken model. If you want that to be loud, have the init service verify the model is listed after pulling rather than trusting the exit code.

Running it in CI

Cache the model volume between runs, keyed on the model name and tag. Without a cache the pull dominates the job and people start skipping the suite on pull requests.
Run with docker compose -f compose.test.yaml up --abort-on-container-exit --exit-code-from app, so the job’s exit status is the test process’s exit status rather than the compose command’s.
Set an explicit CPU limit on the llm service if your runner is shared. An inference server that saturates every core makes the unrelated tests running beside it look slow and flaky.
Fail the build if any test reached the public internet. The simplest enforcement is to run the app service on an internal network with no gateway; a stray real API call then fails loudly instead of quietly costing money.