Test Context Cancellation and Timeouts in Go Code Properly

A function that takes a context.Context makes a promise: when the context ends, I stop. Nothing enforces the promise, and breaking it is easy. A loop that never checks ctx.Done(), a call to a library that ignores contexts, a time.Sleep where a select was needed. The bug does not show in normal use, because contexts usually do not end. It shows at shutdown, when requests that should have stopped keep running, or when a client times out and the server keeps working for nobody.

So the promise needs a test. It is a short one, with three parts: control the context, make the end happen, and assert on both the error and the speed.

The code

package main

import (
	"context"
	"errors"
	"time"
)

// Slow stands for a call that takes d unless the context ends first.
func Slow(ctx context.Context, d time.Duration) error {
	select {
	case <-time.After(d):
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// WithRetry retries Slow until the context ends.
func WithRetry(ctx context.Context, attempts int) (int, error) {
	for i := 1; i <= attempts; i++ {
		if err := Slow(ctx, 10*time.Millisecond); err == nil {
			return i, nil
		}
		if errors.Is(ctx.Err(), context.Canceled) {
			return i, ctx.Err()
		}
	}
	return attempts, errors.New("gave up")
}

func main() {}

Slow stands for any operation that takes time: a request, a query, a batch. It waits for a timer or for the context, whichever comes first, and returns ctx.Err() when the context wins. WithRetry calls it up to a number of times and gives up early if the context was cancelled. Both are small, and both can be wrong in obvious ways.

Four cases

package main

import (
	"context"
	"errors"
	"testing"
	"time"
)

func TestDeadlineExceeded(t *testing.T) {
	ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond)
	defer cancel()
	start := time.Now()
	err := Slow(ctx, time.Hour)
	t.Logf("err=%v isDeadline=%v returnedQuickly=%v", err, errors.Is(err, context.DeadlineExceeded), time.Since(start) < time.Second)
}

func TestCancelWhileWaiting(t *testing.T) {
	ctx, cancel := context.WithCancel(t.Context())
	errc := make(chan error, 1)
	go func() { errc <- Slow(ctx, time.Hour) }()
	cancel()
	select {
	case err := <-errc:
		t.Logf("err=%v", err)
	case <-time.After(2 * time.Second):
		t.Fatal("did not stop after cancel")
	}
}

func TestCancelCause(t *testing.T) {
	ctx, cancel := context.WithCancelCause(t.Context())
	cancel(errors.New("shutting down"))
	t.Logf("Err=%v Cause=%v", ctx.Err(), context.Cause(ctx))
}

func TestRetryStopsOnCancel(t *testing.T) {
	ctx, cancel := context.WithCancel(t.Context())
	cancel()
	n, err := WithRetry(ctx, 5)
	t.Logf("attempts=%d err=%v", n, err)
}

go test -count=1 -v . 2>&1 | grep -vE "^(=== RUN|=== PAUSE|=== CONT|PASS|ok)" | sed -E "s/ \([0-9.]+s\)//; s#^ +[a-z_]+\.go:[0-9]+: #    #"

    err=context deadline exceeded isDeadline=true returnedQuickly=true
--- PASS: TestDeadlineExceeded
    err=context canceled
--- PASS: TestCancelWhileWaiting
    Err=context canceled Cause=shutting down
--- PASS: TestCancelCause
    attempts=1 err=context canceled
--- PASS: TestRetryStopsOnCancel

A deadline. TestDeadlineExceeded gives Slow a context that expires in 20 ms and asks it to wait an hour. It returned context deadline exceeded, and errors.Is(err, context.DeadlineExceeded) was true. The test also checks that the call returned in under a second. That time check is the assertion that matters: a function that ignored the context would wait the full hour, and the test would hang. If you only check the error, a function that waited out the hour and then returned the context error would pass.

Use errors.Is, not ==. A function that wraps the error (fmt.Errorf("fetch: %w", err)) is still correct, and == would reject it. The same applies to context.Canceled.

An external cancel. TestCancelWhileWaiting starts Slow in a goroutine, cancels the context from the test, and waits on a channel for the result with a two second guard. The error is context canceled. The guard is the second half of the pattern: the select with time.After(2 * time.Second) turns a function that ignores cancellation into a test failure with a message, not a stuck process. Without it, the test would hit the global -timeout after ten minutes and print a goroutine dump, which is accurate but slow. Keep the guard much longer than the code should need, so a loaded machine never trips it, and short enough that a real bug is reported quickly.

A reason for cancelling. Since Go 1.20, context.WithCancelCause lets the canceller say why. TestCancelCause cancels with errors.New("shutting down") and reads Err=context canceled Cause=shutting down. ctx.Err() stays context.Canceled so that existing checks keep working, and context.Cause(ctx) returns the specific reason. If your code distinguishes a shutdown from a client disconnect, test that the cause reaches the logs or the error it returns.

A context that is already dead. TestRetryStopsOnCancel cancels first and then calls WithRetry. The result is attempts=1 err=context canceled: one attempt, then no more. A retry loop that does not check the context between attempts keeps trying after the caller has gone. That is the case most likely to be wrong in real retry code, and the cheapest to test, because the context is cancelled before the call and no timing is involved. The retry testing article goes further with backoff.

Use t.Context() as the parent

Go 1.24 added t.Context(), a context that is cancelled just before the test's cleanup functions run. Using it as the parent means that any goroutine the test starts and forgets is told to stop when the test ends, instead of running into the next test. All four tests derive from it. The older equivalent is context.WithCancel(context.Background()) plus t.Cleanup(cancel), which is still right for code that must compile on earlier versions.

Timing, without being flaky

Tests with deadlines have a trap: choosing durations. A 20 ms deadline and a one second bound are fine on any machine. Do not write a test that asserts a call took between 90 and 110 ms, because a loaded runner will fail it. Assert upper bounds, generously, and rely on the deadline to do the stopping.

When the code under test reads the clock itself, as a rate limiter or a cache does, a real deadline is the wrong tool and a virtual clock is the right one: testing/synctest, covered in stopping time-based flakes, runs a one-hour timeout in no time. Use it when a test would otherwise wait for something to time out. For a deadline that you control from the test, a real 20 ms is cheap and clear.

What to check in the code under test

• Every blocking call has a context-aware form: select with ctx.Done(), http.NewRequestWithContext, db.QueryContext, exec.CommandContext. A test that cancels around a call to the standard library proves the context actually reached it.
• Loops that do CPU work check ctx.Err() every so often; a million-iteration loop with no check ignores cancellation until it finishes.
• Goroutines the function starts exit when the context ends. Test that by waiting for them, as in the goroutines and channels article.
• Resources acquired are released on the cancel path: files closed, locks unlocked, connections returned.

A table for the cancel paths

When a package has many functions that take a context, a table keeps the check uniform. Each row names a function and supplies a closure that calls it with the given context. The test loops over the rows, cancels the context first, and asserts that every function returns an error matching context.Canceled within the guard time. Adding a function to the package then means adding a row, and forgetting to is visible in review. It also finds the function that quietly started ignoring its context after a refactor, which is how most of these bugs arrive.

Limits

A test with a cancelled context shows that the code stops; it does not show that the work done before stopping was safe. A half-written file, a transaction left open, a message acknowledged but not processed: these are correctness bugs on the cancel path and need their own assertions. The standard documentation for the package is at pkg.go.dev/context, and the Go blog's article on contexts explains the design.