はじめに

TIG真野です。

フューチャー技術ブログ Go 1.24 リリース連載の2本目です（公開は1ヶ月程度遅れたので、時系列では7本目です）。

Go 1.24で追加された testing.Context (T.Context) について紹介します。

func (c *T) Context() context.Context`

Cleanup に登録された関数が呼び出される直前にキャンセルされるコンテキストを返す
Cleanup でテストまたはベンチマークが完了する前に、Context.Done によってシャットダウンされるリソースを待機できる

testing.Context は定期的に欲しいな⇛プロポーサルが却下されているな..を繰り返してきたので、Go 1.24で入ってとても嬉しく思っています。2年半くらい前にもポストしていました。

testing.Context() のヘルパー関数欲しい（context.Background() 書くのがやや冗長に感じる）
— Junki Mano (@ma91n) September 20, 2022

注目度が高いこともあって、 testing.Context について詳細に説明した記事がいくつも公開されています。

そのため、testing: reconsider adding Context method to testing.T #36532 のIssueの議論のやり取りや、仕様について私なりに気になった点を深堀りすることを中心にします。

Discussion #36532サマリ

T.Context() を追加する提案は観測できる範囲で2016年ごろから断続的に行われていた
T.Context() から取得したテスト終了後にキャンセルするコンテキストを利用することで、例えば、テストコード上でgoroutine呼び出ししていた場合も、テストの停止とタイミングを併せて停止させることができると期待された
#16221 では実装まで行ったが、コンテキストをキャンセルして実際にgoroutineが全て停止したことを確認待ちする手段が無かった。そのため、テストコード上で起動したgoroutineが残ったまま次のテストが動く可能性があり、課題が解決しておらず不完全であるためリバートされた（※後述）
2020年3月リリースのGo 1.14 で t.Cleanup() が実装され、これと組み合わせると待ち合わせが可能となり、課題が消えたため T.Context() が追加可能となった。プロポーサル自体は2020年に起票されていて、2024年6月ごろに議論が再び活発になった流れでした
T.Context() が存在しない場合でも、暫定的な回避方法は存在していたが（※後述）、それでもT.Context() が追加したほうが自然であり、また副作用も少ないと見てゴーの判断となった

Try it

Issueの説明に記載されたコードを少し改変し動かします。

tcontext_test.go

package main

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"testing"
	"time"
)
func TestTContext_Wait(t *testing.T) {
	ctx := t.Context()
	var wg sync.WaitGroup

    // t.Cleanupが呼ばれるタイミング（≒テスト終了時）に
    // wg.Waitで全てのgoroutineタスクが終了するまで待ち合わせする
    // ≒ TestFoo()終了後に別のテストが起動しても、TestFooのテストで起動した
    // 全てのgoroutineが終了していることを保証できる
    t.Cleanup(wg.Wait)

	wg.Add(1)

    // 別goroutineで非道実行
    // T.Context()で取得したctxでテスト終了時にキャンセルされる
	go worker(ctx, &wg, "Worker1")

	time.Sleep(500 * time.Millisecond) // しばらく worker を実行
}

func worker(ctx context.Context, wg *sync.WaitGroup, workerName string) {
	defer wg.Done()
	for {
		select {
		case <-ctx.Done():
			fmt.Printf("%s: context canceled: %v\n", workerName, ctx.Err())
			return
		default:
			fmt.Printf("%s: working...\n", workerName)
			time.Sleep(100 * time.Millisecond) // 何かしら重い処理に相当
		}
	}
}

これを実行します。

$ go test -run TestTContext_Wait
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: context canceled: context canceled
PASS
ok      Go124TContext   0.533s

結果は time.Sleep(500 * time.Millisecond) で指定した 500ms だけworker() のループが5回実行され、その後 TestTContext_Wait() が終了。コンテキストがキャンセルされ、Worker1: context canceled: context canceled というメッセージを出してworkerが終了。その後、 t.Cleanup() が呼ばれるという流れです。

図にすると以下にようなイメージです。

Meramid live editor

コンソールログや上図を見ると、 t.Cleanup() の wg.Wait() 不要じゃね？と思うでしょう（私は思いました）。そのため、別のコードを書いてみます。

package main

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"
)

func worker(ctx context.Context, wg *sync.WaitGroup, workerName string) {
	defer wg.Done()
	for {
		select {
		case <-ctx.Done():
			fmt.Printf("%s: context canceled: %v\n", workerName, ctx.Err())
			return
		default:
			fmt.Printf("%s: working...\n", workerName)
			time.Sleep(100 * time.Millisecond) // 何かしら重い処理に相当
		}
	}
}

func TestTContext_Wait(t *testing.T) {
	ctx := t.Context()
	var wg sync.WaitGroup
-	t.Cleanup(wg.Wait)

	wg.Add(1)
	go worker(ctx, &wg, "Worker1")

	time.Sleep(500 * time.Millisecond) // しばらく worker を実行
}

+func TestTContext_Wait_Next(t *testing.T) {
+	fmt.Println("🔜TestTContext_Wait_Next開始")
+	time.Sleep(500 * time.Millisecond)
+	fmt.Println("🔚️TestTContext_Wait_Next終了")
+}

これを実行すると以下のようになります。

$ go test -run TestTContext_Wait
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: working...
🔜TestTContext_Wait_Next開始
Worker1: context canceled: context canceled
🔚️TestTContext_Wait_Next終了
PASS
ok      Go124TContext   1.022s

興味深いのは、 Worker1: context canceled: context canceled が、次のテストである TestTContext_Wait_Next が起動された後に呼び出されていることです。つまり、TestTContext_Wait で呼び出されたgoroutineが、次のテスト中に存命であるという状態が起こりえるということです。

サマリのタイミングでさらっと説明した通り、 T.Cleanup() が存在しないタイミングでは、 T.Context() が存在しても中途半端だというのは、こういったことを指していました。

一応イメージを図にします。t.Cleanup() で終了を待ち合わせる重要性が分かりました。

Meramid live editor

気になった点

私がリリースノートのT.Context() 部分を読んだところで、気になった点を残していきます。

`t.Fail()` された場合にはキャンセルされるのか？

キャンセルされませんでした（ドキュメント上も、T.Cleanup()前にキャンセルされるとあるので、それはそう）。

func TestTContext_Fail(t *testing.T) {
	ctx := t.Context()
	var wg sync.WaitGroup

	wg.Add(1)
	go worker(ctx, &wg, "Worker1")

	time.Sleep(500 * time.Millisecond) // しばらく worker を実行

	t.Fail() // テストを失敗マークを付ける、テストコードは継続

	// コンテキストがキャンセルされたか確認する
	if errors.Is(ctx.Err(), context.Canceled) {
		fmt.Println("t.Context() was canceled")
	} else {
		fmt.Println("t.Context() was not canceled")
	}
}

実行結果。

$ go test -run TestTContext_Fail
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: working...
t.Context() was not canceled
--- FAIL: TestTContext_Fail (0.51s)
FAIL
exit status 1
FAIL    Go124TContext   0.511s

テストがタイムアウトした場合に呼ばれるのか？

キャンセルされませんでした。

func TestTContext_Timeout(t *testing.T) {
	ctx := t.Context()
	var wg sync.WaitGroup

	wg.Add(1)
	go worker(ctx, &wg, "Worker4")
	time.Sleep(500 * time.Millisecond)

	wg.Wait()

	fmt.Println("❌️-timeoutオプションでキャンセルされた場合に到達するはずだが、されない")
}

-timeout オプション付きで実行します。

$ go test -run TestTContext_Timeout -timeout 1s
Worker4: working...
Worker4: working...
Worker4: working...
Worker4: working...
Worker4: working...
Worker4: working...
Worker4: working...
Worker4: working...
Worker4: working...
Worker4: working...
panic: test timed out after 1s
        running tests:
                TestTContext_Timeout (1s)

goroutine 9 [running]:
testing.(*M).startAlarm.func1()
(中略)

FAIL    Go124TContext   1.011s

従来通り panic で終了する感じです（それはそう）。考えてみればテスト全体の実行時間のタイムアウトなので、プロセスごと落とすで良く、何かハンドリングしたいケースは少なそうです（DBコネクションを明示的にクローズしたいとかはあるかもですが）。

`TestMain()` で呼べないのか？

func TestMain(m *testing.M) {...} では呼べないです。 testing.M には追加されていないです。

ただし、testing.B には追加されています。

サブテストで `T.ContexT()` を呼んだらどうなるのか？

サブテストのスコープでキャンセルされましｔ

package main

import (
	"fmt"
	"sync"
	"testing"
	"time"
)

func TestSubtests(t *testing.T) {
	// 親テストのCleanup
	t.Cleanup(func() {
		fmt.Println("Parent test cleanup")
	})

	tests := []struct {
		name       string
		workerName string
	}{
		{name: "Subtest1", workerName: "Worker1"},
		{name: "Subtest2", workerName: "Worker2"},
		{name: "Subtest3", workerName: "Worker3"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := t.Context() // サブテストのコンテキストを取得

			var wg sync.WaitGroup
			wg.Add(1)

			go worker(ctx, &wg, tt.workerName)

			t.Cleanup(func() {
				fmt.Printf("%s: cleanup\n", tt.name)
				wg.Wait() // ゴルーチンの終了を待つ
			})

			time.Sleep(300 * time.Millisecond)
			fmt.Printf("%s: finishing\n", tt.name)
		})
	}
}

実行します。サブテストごとに設定したt.Cleanup() の呼ばれる前（ログ上はworkerにキャンセルが伝播されて後ログ出力するまでの遅延があり、入れ帰っていますが）に、キャンセルされていることがわ

$ go test -run TestSubtests
Worker1: working...
Worker1: working...
Worker1: working...
Subtest1: finishing
Subtest1: cleanup
Worker1: context canceled: context canceled
Worker2: working...
Worker2: working...
Worker2: working...
Subtest2: finishing
Subtest2: cleanup
Worker2: context canceled: context canceled
Worker3: working...
Worker3: working...
Worker3: working...
Subtest3: finishing
Subtest3: cleanup
Worker3: context canceled: context canceled
Parent test cleanup
PASS
ok      Go124TContext   0.942s

従来の回避方法は？

T.Context() を用いずとも、context.WithCancel() で同等のことが実現できました。Issueのやりとりでは、ctx, cancel := context.WithCancel(context.Background()); defer cancel() で代用できるし、これらは2行とシンプルだし、新しくAPIを追加する必要はないのではないか？といったコメントもありました。

context.WithCancel() で書いたコードです。

func TestTContext_WorkAround(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup

	t.Cleanup(func() {
		cancel()
		wg.Wait()
		fmt.Println("✅️TestTContext_WorkAround: finished")
	})

	wg.Add(1)
	go worker(ctx, &wg, "Worker1")

	time.Sleep(500 * time.Millisecond) // しばらく worker を実行
}

実行します。

$ go test -run TestTContext_WorkAround
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: working...
Worker1: context canceled: context canceled
✅️TestTContext_WorkAround: finished
PASS
ok      Go124TContext   0.529s

これを見ると、t.Cleanup() でコンテキストをキャンセルし、かつ wg.Wait() でテストコード中に起動したgoroutineが全て終了していることを確認して、テストが終了しています。 T.Context() が無くても大丈夫です。では、なぜ追加が決まったのでしょうか。

これについては Ianさんのコメントが流れを決めた気がしています。主張としては以下です。

背景:

複数の企業でテストを追加した経験上、テスト数が増えるとリソースリークに対する課題が出る（既存のテストと似た別のテストを追加した後、テストを全実行すると動かなくなることが1回は出てくる）
各テストの先頭に defer を含むキャンセル可能なコンテキストに書き換えすることで、回避はできた
ただ、その後にまたキャンセル可能なコンテキストを使っていないテストが紛れ込むことがあり、都度修正した

主張:

（T.Context() や context.WithCancel() を問わず）テストの終了時にコンテキストをキャンセルしないと、リークの問題が発生し、デバッグが困難
上記のような経験が無いと、通常 ctx, cancel := context.WithCancel(context.Background()); defer cancel() を利用することはないし、あったとしても一貫して使用するには面倒

さらに、「context.Background() を T.Context() に書き換えれば、テストの最後にキャンセルされたコンテキストの恩恵を受けるか、最悪の場合でも影響を受けない」とあり、確かに副作用はほぼないと感じました。

context.WithCancel()とdeferや t.Cleanup() の組み合わせで代替できるが、多くの開発者が context.Background()をテストで使ってしまうことで、不具合が出てしまっている。これを T.Context() でそのまま置き換えて使ってしまっても、メリットはあれど副作用は考えにくい、ということで、納得感がありました

さいごに

T.Cleanup() が T.Context() 追加するために重要なピースだったということはIssueを読めばすぐ分かりましたが、具体的にどういう意味なのか、コードを動かしてみるまでピンと来ませんでした。手を動かして疑問点を潰しこむ重要性を、Goリリースノート連載はいつも私に教えてくれます。

超個人的には、sync.WaitGroup や errgroup を使用する機会は少ないためテストでハマることは少ないため、context.Background() って書くより t.Context() で書け、文字数が減るというのが一番の嬉しいポイントでした。

最近、業務でGoを書く機会が減っていますが、引き続きGo情報はウォッチして楽しみながら精進したいと思います。

Go 1.24リリース連載 testing.Context