AI-Native Updated 2026-09 11 min read View as Markdown

Testing Go that calls a language model

Q: How do I test streaming?

Have your httptest handler write chunks with an explicit Flush() between them, and vary the chunk boundaries — including one that splits a multi-byte UTF-8 character. That, plus a handler that closes the connection halfway, covers most streaming bugs.

Go is increasingly the language the AI infrastructure is written in rather than the AI feature. That changes what you test — and interfaces plus httptest make it unusually pleasant.

Go turns up in this domain in a particular role: the gateway, the proxy, the worker, the orchestrator. The service that fans out to three providers, enforces a budget, retries, streams back, and records what happened.

That shape is good news for testing, because the interesting behaviour is concurrency, cancellation and failure handling — which is exactly what Go's standard library is best at exercising.

The interface is the whole design#

Define it in the package that uses it, keep it small, and everything above becomes testable with no network.

internal/llm/llm.go

package llm

type Request struct {
    System   string
    Messages []Message
    MaxTokens int
}

type Response struct {
    Text  string
    Usage Usage
}

// Client is what the rest of the application depends on.
type Client interface {
    Complete(ctx context.Context, req Request) (Response, error)
    Stream(ctx context.Context, req Request) (iter.Seq2[string, error], error)
}

Two methods. A real implementation, a fake, and — as it turns out — you rarely need the fake, because httptest gives you something better.

Prefer httptest to a hand-written fake#

A fake Client tests your code against your assumptions. An httptest.Server tests it against real HTTP semantics: status codes, headers, chunked transfer, connections that close mid-body, contexts that cancel.

internal/llm/client_test.go

func TestRetriesOn429ThenSucceeds(t *testing.T) {
    var calls atomic.Int32
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if calls.Add(1) == 1 {
            w.Header().Set("Retry-After", "0")
            w.WriteHeader(http.StatusTooManyRequests)
            return
        }
        _ = json.NewEncoder(w).Encode(apiResponse{Text: `{"category":"billing"}`})
    }))
    defer srv.Close()

    c := New(srv.URL, WithMaxRetries(3))
    got, err := c.Complete(t.Context(), Request{System: "s"})
    if err != nil {
        t.Fatalf("Complete() error = %v", err)
    }
    if calls.Load() != 2 {
        t.Errorf("calls = %d, want 2", calls.Load())
    }
    if got.Text == "" {
        t.Error("empty text")
    }
}

The cases worth writing, all of which are real production failures and none of which need a model:

Test	What it catches
429 with `Retry-After`	retry logic, and that you honour the header
500 on every attempt	that you give up rather than retry forever
Response body closed mid-stream	partial-write handling
Server sleeps past the deadline	that your context timeout actually fires
Malformed JSON	parser never panics
Empty 200	the "provider returned nothing" path

Cancellation is the test people skip#

func TestClientDisconnectCancelsUpstream(t *testing.T) {
    released := make(chan struct{})
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        <-r.Context().Done()      // upstream sees the cancellation
        close(released)
    }))
    defer srv.Close()

    ctx, cancel := context.WithCancel(t.Context())
    go func() { time.Sleep(20 * time.Millisecond); cancel() }()

    c := New(srv.URL)
    if _, err := c.Complete(ctx, Request{}); !errors.Is(err, context.Canceled) {
        t.Fatalf("err = %v, want context.Canceled", err)
    }
    select {
    case <-released:
    case <-time.After(time.Second):
        t.Fatal("upstream request was never cancelled — you are still paying for this")
    }
}

That test pays for itself the first time someone closes a browser tab.

And because this is Go, put the leak check in place once and forget it:

func TestMain(m *testing.M) { goleak.VerifyTestMain(m) }

Streaming code with a for range over a channel is the most common source of goroutine leaks in this kind of service, and goleak turns each one into a test failure rather than a slow memory climb.

Contract tests, table-driven#

Test what your code does with a response, not what the model says.

func TestParseCategory(t *testing.T) {
    tests := []struct {
        name string
        raw  string
        want Category
        ok   bool
    }{
        {"plain", `{"category":"billing"}`, Billing, true},
        {"fenced", "```json\n{\"category\":\"billing\"}\n```", Billing, true},
        {"preamble", "Sure!\n{\"category\":\"billing\"}", Billing, true},
        {"wrong case", `{"category":"Billing"}`, Billing, true},
        {"unknown value", `{"category":"refunds_and_returns"}`, "", false},
        {"misspelled key", `{"categorie":"billing"}`, "", false},
        {"prose", "I think this is a billing issue.", "", false},
        {"truncated", `{"category":"billing"`, "", false},
        {"empty", ``, "", false},
    }
    for _, tt := range tests {
        t.Run(tt.name, func(t *testing.T) {
            got, ok := ParseCategory(tt.raw)
            if ok != tt.ok || got != tt.want {
                t.Errorf("ParseCategory() = %q, %v; want %q, %v", got, ok, tt.want, tt.ok)
            }
        })
    }
}

The contract is never panic, never invent. Returning a plausible wrong category is worse than returning false.

Golden files instead of cassettes#

Go's idiom for recorded responses is a golden file plus an -update flag. Simpler than the equivalents elsewhere and it lives in the standard toolchain.

var update = flag.Bool("update", false, "rewrite golden files")

func TestPromptRendering(t *testing.T) {
    got := BuildPrompt(fixtureQuery(), fixtureChunks())
    golden := filepath.Join("testdata", "prompt.golden")

    if *update {
        if err := os.WriteFile(golden, []byte(got), 0o644); err != nil {
            t.Fatal(err)
        }
    }
    want, err := os.ReadFile(golden)
    if err != nil {
        t.Fatal(err)
    }
    if diff := cmp.Diff(string(want), got); diff != "" {
        t.Errorf("prompt mismatch (-want +got):\n%s", diff)
    }
}

shell

go test ./... -update       # re-record, then read the diff before committing

Reading that diff is the point. A prompt change you did not intend shows up here before it shows up in your eval scores.

Store recorded API responses in testdata/ the same way and serve them from your httptest handler. Scrub credentials before committing, and re-record on a schedule — a golden file from eighteen months ago is testing a model that no longer exists.

Concurrency: the part Go actually adds#

If your service fans out, test the fan-out. This is where a Go LLM gateway earns its keep and where it breaks.

func TestFanOutRespectsConcurrencyLimit(t *testing.T) {
    var inFlight, peak atomic.Int32
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        n := inFlight.Add(1)
        for {
            p := peak.Load()
            if n <= p || peak.CompareAndSwap(p, n) {
                break
            }
        }
        time.Sleep(10 * time.Millisecond)
        inFlight.Add(-1)
        _, _ = w.Write([]byte(`{"text":"ok"}`))
    }))
    defer srv.Close()

    if err := ClassifyAll(t.Context(), New(srv.URL), makeInputs(50), WithLimit(5)); err != nil {
        t.Fatal(err)
    }
    if got := peak.Load(); got > 5 {
        t.Errorf("peak concurrency = %d, want <= 5", got)
    }
}

Run it with -race, always. Unbounded fan-out at a rate-limited provider is the most common way one of these services falls over, and it is a five-line errgroup.SetLimit fix once you know.

Evals#

Nightly, behind a build tag so they never run in the normal suite.

//go:build eval

package eval

func TestClassificationAccuracy(t *testing.T) {
    cases := loadDataset(t, "testdata/dataset.jsonl")
    client := llm.NewFromEnv()

    var hits int
    for _, c := range cases {
        got, err := Classify(t.Context(), client, c.Input)
        if err != nil {
            t.Errorf("%q: %v", c.Input, err)
            continue
        }
        if got == c.Expect {
            hits++
            continue
        }
        t.Logf("MISS %q: got %s, want %s", c.Input, got, c.Expect)
    }

    acc := float64(hits) / float64(len(cases))
    t.Logf("accuracy %.1f%% on %d cases", acc*100, len(cases))
    if acc < 0.90 {
        t.Errorf("accuracy %.1f%% below threshold 90%%", acc*100)
    }
}

shell

go test -tags eval ./eval/...      # nightly only

t.Logf for the misses rather than t.Errorf — the individual failures are information, the aggregate is the gate. Read the misses; that list is where the next prompt change comes from.

The whole loop#

makefile

test:                                    # every edit, no network, no tokens
	go test -race -count=1 -timeout 60s ./...

eval:                                    # nightly
	go test -tags eval -count=1 -timeout 20m ./eval/...

Fast, honest, free, and it exercises the parts of an LLM service that actually break.

Common questions#

Should I use an SDK or call the HTTP API directly?#

For a gateway or proxy — the shape Go usually takes here — the raw API is often simpler. You are already handling streaming, retries and budgets yourself, and an SDK mostly adds a layer between you and the behaviour you are trying to control. For an application that makes occasional calls, use the SDK.

Why `httptest` rather than a mock client?#

A mock tests your code against your assumptions about HTTP. httptest tests it against real HTTP: connection resets, chunked encoding, contexts that cancel mid-body. Those are where the bugs are, and they are the cases a hand-written fake never simulates.

How do I test streaming?#

Have your httptest handler write chunks with an explicit Flush() between them, and vary the chunk boundaries — including one that splits a multi-byte UTF-8 character. That, plus a handler that closes the connection halfway, covers most streaming bugs.

How do I stop the eval run costing a fortune?#

Build tags keep it out of the normal suite, which is most of it. Beyond that: bound the fan-out with errgroup.SetLimit, use the batch API where latency does not matter, and enforce a budget in the context — see tracking and cutting token costs in Go.

Do golden files not just get blindly updated?#

They do if nobody reads the diff, which is why the discipline is to run -update and then read git diff testdata/ before committing. That habit is the entire value; without it a golden file is a test that always passes.

Get the Go agent pack

A battle-tested AGENTS.md, the review checklist, and the failure-mode cheat sheet for Go. One email, then occasional updates when the tooling shifts. No course pitch.

Unsubscribe in one click. We never sell the list. Or just take the AGENTS.md now — no email needed.