channelを使うテストで無限待ちしないようにする

Goで並列に動くものをテストしたいとき、順序を一定に保つためにchannelをよく使う。

しかし、実装を間違えるといつまでもchannelの送受信が終わらずに無限に待ち続けてテストが止まることがある。あちこちでchannelの送受信をしていると、どこで止まったのかが一見ではわかりにくい。

そこで、channelの送受信にタイムアウトを設定して一定以上の時間送受信できなかった場合エラーにしてしまうと良い。

// channelから値を受信する（タイムアウト付き）
ch := make(chan string)

select {
case val := <-ch:
    log.Printf("received: %v", val)
case <-time.After(1*time.Second):
    log.Printf("timed out")
}

確かにこれでいいのだが、テストで何箇所もchannelの送受信を入れていると、そのたびに select ブロックを書くのはコードが長くなり気持ち悪い。そこで、タイムアウト付きの channel 送受信をテスト用補助関数として共通化した。

func MustSendChan(t *testing.T, channel, value interface{}, timeout time.Duration) {
    t.Helper()

    cases := []reflect.SelectCase{
        {Dir: reflect.SelectSend, Chan: reflect.ValueOf(channel), Send: reflect.ValueOf(value)},
        {Dir: reflect.SelectRecv, Chan: reflect.ValueOf(time.After(timeout))},
    }
    chosen, _, _ := reflect.Select(cases)
    if chosen == 1 { // index 1: timeout
        t.Fatalf("channel send timed out")
    }
}

func MustReceiveChan(t *testing.T, channel interface{}, timeout time.Duration) interface{} {
    t.Helper()

    cases := []reflect.SelectCase{
        {Dir: reflect.SelectRecv, Chan: reflect.ValueOf(channel)},
        {Dir: reflect.SelectRecv, Chan: reflect.ValueOf(time.After(timeout))},
    }
    chosen, value, _ := reflect.Select(cases)
    switch chosen {
    case 0: // received from channel
        return value.Interface()
    case 1: // timeout
        t.Fatalf("channel receive timed out")
    }
    // channel closed
    return nil
}

ch := make(chan string)

// channelに値を送信する、1秒以上送信できなかったらタイムアウト
MustSendChan(t, ch, "value", 1*time.Second)

// channelから値を受信する、1秒以上受信できなかったらタイムアウト
// interface{} 型で受信されるので、適宜castしてあげる必要がある
value := MustReceiveChan(t, ch, 1*time.Second).(string)

Goのchannelは何の型を送受信するかで型が異なるので、あらゆるchannelを受け入れるために interface{} 型を使っている。そして、reflect.Selectを使うことで reflect.Value から動的に select-case ブロックと同等のことができる。

また、t.Helperを挿入してテスト失敗時のスタックトレースにこの関数は含まないようにした。

これでテスト内のchannel送受信がスッキリしつつ、実装ミスの原因にも気づきやすくなった！ただ、共通化した結果 interface{} を使ってしまったので、こういうときにGenericsが欲しくなるんだなぁ・・と実感した。