Goのtesting.TにContext()メソッドが追加されそうという話

この記事はsyumai Advent Calendar 2024の3日目の記事です。

火曜なので、本日は「Go」がテーマとなります。

(*testing.T).Context()とは

(*testing.T).Context() は、以下のProposalでAcceptされた、Goのテストで使用できる新しいメソッドです。

github.com

このメソッドから返却されるContextは、t.Cleanup() によって登録されたクリーンアップ関数が呼び出される直前にキャンセルされます。

Proposalの最終的なAPIは以下の通りです。

// Context returns a context that's cancelled just before
// [T.Cleanup]-registered functions are called.
//
// Cleanup functions can wait for any resources
// that shut down on Context.Done before the test completes.
func (t *T) Context() context.Context

実装は既にmaster branchにマージされていて、現在Go 1.24のマイルストーンでのリリースが見込まれています。 testing.Tへの追加と言いましたが、実際には、testing.T / B / Fで共通のインタフェースであるtesting.TBへ追加が行われたようです。

使い方

ProposalのDescriptionに記載されているのは、次のような使い方です。

func TestFoo(t *testing.T) {
    // 1. Contextを取得
    ctx := t.Context()
    // 2. WaitGroupを作成
    var wg sync.WaitGroup
    // 3. テストの完了時のクリーンアップ処理でwg.Waitを呼ぶように指定
    t.Cleanup(wg.Wait)
    // 4. WaitGroupにテスト内で使用する処理を登録
    wg.Add(1)
    go func() {
        defer wg.Done()
        // 5. テストのクリーンアップ処理の直前まで待つContextを渡す
        doSomething(ctx)
    }()
}

このコードを読むだけではあまりイメージしにくいかもしれませんが、本機能のメリットは「テストの完了まで待ってくれるContextが入手できる」点にあります。

上記の doSomething 関数をもっと具体化してみましょう。

HTTPサーバーの終了の待機

例えば、main関数で以下のような関数によってHTTPサーバーを起動しているとします（簡単な例という点を踏まえて読んでいただけるとありがたいです）。

このHTTPサーバーは、外から渡されたContextの終了でもってShutdown関数の呼び出しを行い、その終了まで待ちます。

func runServer(ctx context.Context, port string) error {
    h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.Write([]byte("Hello, World!"))
    })
    srv := &http.Server{Addr: fmt.Sprintf(":%s", port), Handler: h}
    go func() {
        if err := srv.ListenAndServe(); err != nil {
            log.Printf("error on serve: %v", err)
        }
    }()

    // 外から渡されたContextの完了を待ってサーバーをShutdownする
    <-ctx.Done()
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    return srv.Shutdown(ctx)
}

テストコードとして、このHTTPサーバーに対してリクエストを行い、その結果を検証することを考えてみます。

このテストは、以下のような単純なGETリクエストで行えるでしょう。

func Test_runServer(t *testing.T) {
    // Contextを用意
    ctx := context.Background()

    // HTTPサーバーを実行 (簡略化のため、WaitGroupをチャネルに置き換えました)
    const testPort = "12345"
    doneCh := make(chan struct{})
    wait := func() { <-doneCh }
    // HTTPサーバーが完全に終了するまで待つ
    t.Cleanup(wait)
    go func() {
        defer close(doneCh)
        runServer(ctx, testPort)
    }()

    // GETリクエストを実行
    res, err := http.Get(fmt.Sprintf("http://localhost:%s", testPort))
    if err != nil {
        t.Fatal(err)
    }
    defer res.Body.Close()

    // GETリクエストの結果を全て読み取る
    body, err := io.ReadAll(res.Body)
    if err != nil {
        t.Fatal(err)
    }

    // 結果を比較
    want := "Hello, World!"
    got := string(body)
    if got != want {
        t.Fatalf("want: %q, got: %q", want, got)
    }
}

実は、ここではあえて誤った例を示しました。問題は、テスト関数の一番上で用意しているContextです。

ctx := context.Background()

単に context.Background() を使用すると、HTTPサーバーが終了することがなく、クリーンアップの wait 関数の呼び出しで待ち続けてしまいます。

この問題は、 context.WithCancel を使うことで解消できます。

func Test_runServer(t *testing.T) {
    // Contextを用意
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    ...
}

このようにすると、テスト全体が完了した後にContextがキャンセルされ、その通知がHTTPサーバーに届き、無事シャットダウン処理が行われます。

しかしながら、ここで更にサブテストが存在したらどうなるでしょうか？

例として、先ほどのGETリクエストによる検証を10並列にしてみます。

func Test_runServer(t *testing.T) {
    ...
    // 10並列でGETリクエストを実行
    for i := range 10 {
        t.Run(fmt.Sprintf("test %d", i), func(t *testing.T) {
            t.Parallel()

            res, err := http.Get(fmt.Sprintf("http://localhost:%s", testPort))
            if err != nil {
                t.Fatal(err)
            }
            defer res.Body.Close()
            ...
            want := "Hello, World!"
            got := string(body)
            if got != want {
                t.Fatalf("want: %q, got: %q", want, got)
            }
        })
    }
}

このようにすると、見事全てのテストケースが失敗します。

    main_test.go:32: Get "http://localhost:12345": dial tcp [::1]:12345: connect: connection refused
    main_test.go:32: Get "http://localhost:12345": dial tcp [::1]:12345: connect: connection refused
    main_test.go:32: Get "http://localhost:12345": dial tcp [::1]:12345: connect: connection refused
    ...
--- FAIL: Test_runServer (0.00s)
    --- FAIL: Test_runServer/test_2 (0.00s)
    --- FAIL: Test_runServer/test_9 (0.00s)
    --- FAIL: Test_runServer/test_3 (0.00s)
    ...
FAIL

このようなテストでは、Contextのキャンセル処理は defer ではなく t.Cleanup で行う必要があります。

注意すべきなのは、 t.Cleanup に登録された処理の呼び出しがLIFO (後入れ先出し) となる点です。つまり、（当然ではありますが）シャットダウンの待機処理より前にContextがキャンセルされる必要があるので、コード上での呼び出し位置として t.Cleanup(wait) よりも t.Cleanup(cancel) が後ろである必要があります。

func Test_runServer(t *testing.T) {
    ctx, cancel := context.WithCancel(context.Background())
    // ここで defer cancel() しないし、 t.Cleanup(cancel) もしない

    const testPort = "12345"
    doneCh := make(chan struct{})
    wait := func() { <-doneCh }
    t.Cleanup(wait)
    // この位置でキャンセル処理を登録する。ここより前にすると、テストが永遠に終了しない
    t.Cleanup(cancel)

    go func() {
        defer close(doneCh)
        runServer(ctx, testPort)
    }()
    ...
}

最終形のコードはこちらです。

t.Context() があると、この面倒なContextのキャンセル処理の位置を気にしなくてよくなります。

func Test_runServer(t *testing.T) {
    const testPort = "12345"
    doneCh := make(chan struct{})
    wait := func() { <-doneCh }
    t.Cleanup(wait)
    go func() {
        defer close(doneCh)
        // これだけで、クリーンアップ処理前にContextがキャンセルされることが保証される
        runServer(t.Context(), testPort)
    }()
    ....
}

少々長くなってしまいましたが、この機能が欲しくなるモチベーションは何となくご理解いただけたのではないかと思います。

Proposalの経緯

実は、この提案はこれまでに何度も行われ、その全てが最終的には承認されなかったようです。

• testing: add T.Context() accessor · Issue #16221 · golang/go · GitHub
• Proposal: add `context.Context` objects to `testing.T` EDIT: and `testing.M` · Issue #18182 · golang/go · GitHub
• testing: T.Context usability concerns · Issue #18199 · golang/go · GitHub
• proposal: testing: add T.Context · Issue #18368 · golang/go · GitHub

Issue: 16221 については一度は承認され、実装が行われたものの、後にリバートされたとのことです。

今回のProposalを提出したrogpeppeさんによると、過去のProposalが却下されたのは、Contextによるシャットダウン指示までは行えるが、その完了を待ち受ける方法が存在しなかったためと認識しているようで、「t.Cleanupが実装された今なら改めて検討できるのでは」ということで提案され、承認に繋がったようです。

現在の状況

まだ議論されていないのでどうなるか不透明ですが、t.Context() で得られるContextに更に付加情報を与えられないかというProposalが出ているようです。

github.com

一つはDeadlineです。 t.Deadline() は、テスト実行時に -timeout フラグで指定された期限を値として返しますが、これを t.Context() の返すContextのDeadlineに設定できないかという提案のようです。

もう一つは、trace.Task の登録です。テスト名を trace.Task に登録してContextに紐づけることで、 デバッグ用のログ出力に使えて便利ということのように見えました（詳しくないので曖昧ですみません…） もし詳しい方がいらしたらXなどでコメントいただけるとありがたいです。

こちらのProposalは、t.Context() の追加の決定そのものに影響するものではなさそうなので、このまま行けば恐らくGo 1.24でリリースされるのではないかと思います。

APIテストなどでの待ち受け処理で活用できそうなので、リリースが楽しみですね…！