Plan 9とGo言語のブログ

主にPlan 9やGo言語の日々気づいたことを書きます。

Goでモンキーパッチするライブラリを作った

Go

Goで単体テストを実装する場合、動的な言語のように「テスト実行中に外部への依存を置き換える」といったことはできません。代わりに、

のように、テスト対象をテスト可能な実装に変更しておき、テストの時は外部への依存をモック等に置き換えて実行する場合が多いのではないかと思います。

個人的な体験でいえば、テスト可能な実装に置き換えていく過程で設計が洗練されていく*1ことは度々あるので、面倒を強制されているというよりは設計を整理するための道具といった捉え方をしているのですが、そうは言っても動的な言語に比べると面倒だなと感じるときは少なからずあります。既存の実装がテスト可能になっておらず、変更するコストが高い場合は特にそうですね。

そんなとき、気軽にモンキーパッチできると嬉しいんじゃないかと思って、テストの時だけ関数を置き換えられるようなライブラリを作りました。

github.com

このライブラリはtenntenn/testtimeにとても影響を受けています。

使い方

試しに標準ライブラリの time.Now を置き換えます。具体的なコードは次のようになります。

import (
    "testing"
    "time"

    "github.lufia/plug"
)

func isLeap() bool {
    now := time.Now()
    return (now.Year() % 4) == 0 // 主題ではないのでうるう年の実装は省略
}

func TestIsLeap(t *testing.T) {
    scope := plug.CurrentScopeFor(t)
    key := plug.Func("time.Now", time.Now)
    plug.Set(scope, key, func() time.Time {
        return time.Date(2024, 5, 10, 11, 0, 0, 0, time.UTC)
    })
    if !isLeap() {
        t.Errorf("2024 is a leap year")
    }
}

plug.Func の第1引数は関数の名前を指定します。理由は後述しますが、これは必ず以下の書式で記述してください。

  • (package-path).(function-name)
  • (package-path).(type-name).(method-name)

例を挙げると math/rand/v2.N や net/http.Client.Do などです。標準の go doc が受け取る引数と似せていますが、パッケージ名の省略はできません*2。

これで、TestIsLeap の中で実行した time.Now は固定で2024年5月10日の時刻を返すようになります。スタックを抜けない限り影響は続くので、isLeap 関数が呼び出す time.Time も固定の値を返します。

テストの実行

テストを実行するときは以下のように実行してください。-overlayオプションが必要です。

go test -overlay <(go run github.com/lufia/plug/cmd/plug@latest)

# 分けて書いてもいい
go run github.com/lufia/plug/cmd/plug@latest >overlay.json
go test -overlay overlay.json

-overlay オプションの詳細は、上で挙げたtenntennさんの記事を読んでもらうと良いのですが、ここでは以下のようなことを実行しています。

  • カレントディレクトリのソースコードから plug.Func を探す
  • plug.Func の第2引数を動的に置き換えできるように書き換える
  • 実行スタックに関連づいたスコープを抜けるまで、plug.Set の第3引数に渡す関数で time.Now を置き換える
  • time.Now を呼び出したとき、実行スタックを遡って直近の time.Now を呼び出し、結果を返す
  • 該当する関数が実行スタック上で置き換えられてなければ本物の結果を返す

plug@latest はカレントディレクトリに plug/ というディレクトリを作成しますが、これは実行するたびに生成するので、不要になったら消しても問題ありません。

サブテストで部分的に置き換える

一部のサブテスト実行中だけ、別の値に置き換えたい場合は、サブテストで同じように書くと実現できます。

func TestIsLeap(t *testing.T) {
    scope := plug.CurrentScopeFor(t)
    key := plug.Func("time.Now", time.Now)
    plug.Set(scope, key, func() time.Time {
        return time.Date(2024, 5, 10, 11, 0, 0, 0, time.UTC)
    })
    t.Run("サブテスト", func(t *testing.T) {
        scope := plug.CurrentScopeFor(t)
        plug.Set(scope, key, func() time.Time { ... })
        // これ以降、サブテストの中では別の値を返す
    })
    // サブテストの外では2024年5月の時刻を返す
}

メソッドを置き換える

メソッドも置き換えできます。以下の例では、net/http.Client の Do メソッドを置き換えているので、http.Get にも影響しています。

func TestHTTPClientGet(t *testing.T) {
    scope := plug.CurrentScopeFor(t)
    key := plug.Func("net/http.Client.Do", (*http.Client)(nil).Do)
    plug.Set(scope, key, func(req *http.Request) (*http.Response, error) {
        return &http.Response{StatusCode: 200}, nil
    })
    resp, _ := http.Get("https://example.com")
}

ジェネリック関数を置き換える

型パラメータのある関数は、型ごとに関数を渡します。

func TestMathRand(t *testing.T) {
    scope := plug.CurrentScopeFor(t)
    key := plug.Func("math/rand/v2.N", rand.N[int])
    plug.Set(scope, key, func(n int) int {
        return 3
    })
    fmt.Println(rand.N[int](10))
}

このとき、 rand.N[int] は plug.Set で差し替わった関数が使われますが、 rand.N[int64] は登録していないので本物の実装が使われます。

関数の引数や呼び出し回数を検査する

内部的に呼ばれた回数を持っているので、それを使って期待した通りに呼ばれているかを検査できます。plug.FuncRecorder[T] に渡す構造体のフィールドは、関数引数の名前に対応したものが使われます。このとき、関数引数の名前をブランク指定子(_)にしていると無視します。

func TestRecorder(t *testing.T) {
    scope := plug.CurrentScopeFor(t)
    key := plug.Func("os.Getenv", func(string) string {
        return "dummy"
    })
    var r plug.FuncRecorder[struct {
        Key string `plug:"key"`
    }]
    plug.Set(scope, key, fake).SetRecorder(&r)

    os.Getenv("PATH")
    if r.Count() != 1 {
        t.Errorf("Count = %d; want 1", r.Count())
    }
    if r.At(0).Key != "PATH" {
        t.Errorf("At(0).Key = %s; want PATH", r.At(0).Key)
    }
}

今後の予定

実行のたびに静的解析をして必要なファイルを生成しているので、パッケージが多くなってくると有意に遅くなります。Goツールチェーンとパッケージのバージョンが変わらなければ基本的には生成するファイルも同じものになるので、うまく最適化ができるといいですね。

他にも、ジェネリック型のメソッドに対応したりとか、go build でも使えるようにしたりなど、色々とやりたいことはあります。

捕捉: なぜ文字列のキーを必要としているか

Goでは関数が同一かどうかを比較することができません。ジェネリックでない関数の場合は reflect.ValueOf(os.Getenv).Pointer() を経由することで比較できますし、Linux/AMD64の場合はだいたい期待通りに動きますが reflect.Value.Pointer のドキュメントには以下のように書かれています。

If v's Kind is Func, the returned pointer is an underlying code pointer, but not necessarily enough to identify a single function uniquely. The only guarantee is that the result is zero if and only if v is a nil func Value.

さらにジェネリック関数では、型パラメータごとに異なる関数ポインタが割り当てられるようで、 reflect.Value.Pointer での比較にも失敗します。

func N[T any](n T) {}

func N1[T any](n T) func(T) {
    return N[T]
}

func N2[T any](n T) func(T) {
    return N[T]
}

func main() {
    fmt.Println(reflect.ValueOf(N[int]).Pointer() == reflect.ValueOf(N[int]).Pointer())   // true
    fmt.Println(reflect.ValueOf(N1[int]).Pointer() == reflect.ValueOf(N[int]).Pointer())  // false
    fmt.Println(reflect.ValueOf(N2[int]).Pointer() == reflect.ValueOf(N2[int]).Pointer()) // true
}

runtime.FuncForPC なども含めて色々と試してみたけれど、Go 1.22時点では良い方法がなかったので、今の形に落ち着きました。

*1:テストコードと同様にドキュメントを書いているときにもよく起きる

*2:go docは http.Client と記述すると推測してくれる