gomockで「呼ばれるべき」と指定した関数の非同期な実行を待ちたいのでテストヘルパーを書いてみた

gomockでデータベース操作などをモックすることはよくあると思います。

そのテスト関数内で「呼ばれるべき（EXPECT()）」と指定した関数に非同期に実行される挙動がある場合、テスト関数の書き方によっては実行終了までにその関数が実行されないことがあり、結果としてテストがエラーになってしまうことがあります。そして、エラーになったりエラーにならなかったりしてくると、Flakyなテストの出来上がりです*1。

（「そもそもテストしやすいように〜すべきだ」というテスト容易性については一旦傍に置いておきます）

ものすごい雑コードですが、例えば以下のような感じのテストです*2。

func TestAsyncExec(t *testing.T) {
    ctrl := gomock.NewController(t)
    t.Cleanup(func() {
        ctrl.Finish()
    })
    m := NewMockQuerier(ctrl)
    m.
        EXPECT().
        Exec("INSERT INTO users (name) VALUES ('alice');").
        Return(nil)

    go func() {
        time.Sleep(time.Duration(rand.Intn(1000)) * time.Millisecond)
        m.Exec("INSERT INTO users (name) VALUES ('alice');")
    }()
}

実際に動くコードはこちらです。

実行すると以下のようなエラーが出ます。goroutine内の処理が実行しきる前にテスト関数の実行が完了するからですね。

=== RUN   TestAsyncExec
    controller.go:269: missing call(s) to *main.MockQuerier.Exec(is equal to INSERT INTO users (name) VALUES ('alice'); (string)) /tmp/sandbox509983514/prog.go:24
    controller.go:269: aborting test due to missing call(s)
--- FAIL: TestAsyncExec (0.00s)
FAIL

Program exited.

「テスト関数の終了をちょっとだけ待ってほしい」だけなんですが、なんとなくの待ち時間を決めて time.Sleep() を使うのもちょっとなあと思ってしまいます。

なので、gomockで「呼ばれるべき」と指定した関数が全て実行されるのを一定間隔で確認するテストヘルパーを書いてみました。

func waitMockCall(t *testing.T, ctrl *gomock.Controller, timeout time.Duration) {
    t.Helper()
    called := false
    done := make(chan interface{})

    go func() {
        time.Sleep(timeout)
        close(done)
    }()

    for {
        select {
        case <-done:
            break
        default:
        }

        var c reflect.Value = reflect.ValueOf(ctrl).Elem()
        var cs reflect.Value = c.FieldByName("expectedCalls").Elem()
        i := cs.FieldByName("expected").MapRange()
        called = true
        for i.Next() {
            calls := i.Value()
            if calls.Len() != 0 {
                called = false
            }
        }
        if called {
            break
        }
        time.Sleep(10 * time.Millisecond)
    }

    if !called {
        t.Log("missed call in time")
    }
}

waitMockCall(t *testing.T, ctrl *gomock.Controller, timeout time.Duration) を、テスト関数の末尾で呼び出しておけば、EXPECT()した関数が全て呼ばれるまで一定期間待つようになります。

実際に動くコードはこちらです。

実行すると以下のように成功するはずです。

=== RUN   TestAsyncExec
--- PASS: TestAsyncExec (0.09s)
PASS

Program exited.

feat: Add AllExpectedCallsSatisfied to Controller by FlorianLoch · Pull Request #627 · golang/mock · GitHub らへんのPull Requestが取り込まれていればよかったのですが、まだだったので reflect パッケージを使って使えそうな値を引っ張り出しています。

で、結局どうするのが正解？

みなさんはどうしていますか？是非知りたいです。

*1:「フレイキーテスト」。最近覚えた言葉です。https://www.publickey1.jp/blog/22/itjenkinsdevops_days_tokyo_2022.html

*2:「なんで突然の goroutine なんだよ！」というのは、簡単にするためにテスト関数内に書いているだけで本当は何か別の関数内で実行されるものとしてください