blog.syfm

徒然なるままに考えていることなどを書いていくブログ

Go の t.Cleanup がとてもべんり

Go Notes Programming

Go 1.14 で testing パッケージに新しく t.Cleanup(func()) や b.Cleanup(func()) が導入されました。
最初は今まで defer を使っていたところを置き換えられるくらいしか良いところがないかな〜と思っていましたが、想像以上に柔軟な使い方ができるので今まで使用したパターンを書いておきます。

Cleanup の特徴

テストランナーは panic ハンドラがあるので、Cleanup は panic が起きたとしても常に呼び出されます。例えば、以下のコードではちゃんと called が出力されます。

func Test_main(t *testing.T) {
    t.Cleanup(func() {
        fmt.Println("called")
    })
    panic("")
}

The Go Playground

別 goroutine で panic した場合や t.Fatal 系が呼ばれた場合でも常に呼び出されます。

追記:
柴田さんからの指摘により、goroutine のスケジューリングを待っていなかったので呼び出されていたことがわかりました。

以下のコードを実行すると t.Cleanup は呼ばれていないことがわかります。

The Go Playground

ドキュメントにもある通り、Cleanup に登録した関数は defer と同様に FILO で呼び出されていきます。以下のコードだと 3、2、1 の順番となります。

func Test_main(t *testing.T) {
    t.Cleanup(func() {
        fmt.Println("1")
    })
    t.Cleanup(func() {
        fmt.Println("2")
    })
    t.Cleanup(func() {
        fmt.Println("3")
    })
}

The Go Playground

Cleanup と defer では defer のほうが先に呼び出されます。

func Test_main(t *testing.T) {
    t.Cleanup(func() {
        fmt.Println("cleanup")
    })
    defer func() {
        fmt.Println("defer")
    }()
}

The Go Playground

また、伝搬はしないので t.Run によってサブテストを作ったとしても、サブテストの終了時に Cleanup が呼ばれることはありません。

役に立ちそうなパターン

環境変数の切り替え

テスト時に動的に環境変数を与えたい場合、今までは以下のようなコードを書いていました。

func setEnv(t *testing.T, k, v string) func() {
    old := os.Getenv(k)
    if err := os.Setenv(k, v); err != nil {
        t.Fatal(err)
    }
    return func() {
        if err := os.Setenv(k, old); err != nil {
            t.Fatal(err)
        }
    }
}

func Test_main(t *testing.T) {
    unsetEnv := setEnv(t, "foo", "bar")
    defer unsetEnv()
}

Cleanup を使うと、unsetEnv の呼び出しを caller にさせずに元に戻すことができます。

func setEnv(t *testing.T, k, v string) func() {
    old := os.Getenv(k)
    if err := os.Setenv(k, v); err != nil {
        t.Fatal(err)
    }
    t.Cleanup(func() {
        if err := os.Setenv(k, old); err != nil {
            t.Fatal(err)
        }
    })
}

func Test_main(t *testing.T) {
    setEnv(t, "foo", "bar")
}

クライアントとサーバの用意

Web API サーバのような、クライアントとサーバが存在する場合は、それらの用意を Cleanup を使って簡潔に書くことができます。 今までだと、例えば HTTP サーバのテストをしたい場合は以下のようなコードを書くことになると思います。

func Test_main(t *testing.T) {
    srv := newServer(t)
    go func() {
        if err := srv.ListenAndServe(); err != http.ErrServerClosed {
            t.Error(err)
        }
    }()

    cli := newClient(t)

    // Test with client.

    ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
    defer cancel()
    if err := srv.Shutdown(ctx); err != nil {
        t.Fatal(err)
    }
}

func newServer(t *testing.T) *http.Server {
    mux := http.NewServeMux()

    // Register handlers.

    return &http.Server{Addr: ":8080", Handler: mux}
}

func newClient(t *testing.T) *http.Client {
    return http.DefaultClient
}

このコードだと、テスト関数である Test_main 内にテストの setup/teardown ロジックが多く入り込んでしまっていて可読性も良くないし、別なテスト関数で同じようなことをやりたいときにこれらのロジックを再び書くことになってしまいます。

t.Cleanup を使うと以下のように caller 側の setup/teardown ロジックを callee に隠蔽することができます。

func Test_main(t *testing.T) {
    cli := newClientAndRunServer(t)

    // Test with client.
}

func newServer(t *testing.T) *http.Server {
    mux := http.NewServeMux()

    // Register handlers.

    return &http.Server{Addr: ":8080", Handler: mux}
}

func newClientAndRunServer(t *testing.T) *http.Client {
    srv := newServer(t)
    go func() {
        if err := srv.ListenAndServe(); err != http.ErrServerClosed {
            t.Error(err)
        }
    }()
    t.Cleanup(func() {
        ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
        defer cancel()
        if err := srv.Shutdown(ctx); err != nil {
            t.Fatal(err)
        }
    })

    return http.DefaultClient
}

サーバの起動から停止まで全てを newClientAndRunServer に入れることができるので caller 側はサーバについて何も考えなくても良くなりました。

setup で複数の goroutine を起動したい場合

上記のクライアント・サーバの応用系として、サーバだけでなく非同期で動くワーカー goroutine などが必要である場合も sync.WaitGroup や errgroup.Group を使って簡潔に書くことができます。

func Test_main(t *testing.T) {
    cli := newClientAndRunServer(t)

    // Test with client.
}

func newServer(t *testing.T) *http.Server {
    mux := http.NewServeMux()

    // Register handlers.

    return &http.Server{Addr: ":8080", Handler: mux}
}

func newClientAndRunServer(t *testing.T) *http.Client {
    ctx, cancel := context.WithCancel(context.Background())
    eg, cctx := errgroup.WithContext(ctx)
    t.Cleanup(func() {
        cancel()
        if err := eg.Wait(); err != nil {
            t.Fatal(err)
        }
    })

    srv := newServer(t)
    eg.Go(func() error {
        if err := srv.ListenAndServe(); err != http.ErrServerClosed {
            return err
        }
        return nil
    })
    t.Cleanup(func() {
        ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
        defer cancel()
        if err := srv.Shutdown(ctx); err != nil {
            t.Fatal(err)
        }
    })
    eg.Go(func() error {
        return runWorker1(cctx)
    })
    eg.Go(func() error {
        return runWorker2(cctx)
    })

    return http.DefaultClient
}

べんりですね。