Goメモ-466 (iotestパッケージのメモ)(03-TimeoutReader)

概要

以下、自分用のメモです。忘れないうちにメモメモ。。。

以下の書籍で知ったのですが、testing/iotest というパッケージの存在を知りました。こんなのあったんですね。

Go言語 100Tips ありがちなミスを把握し、実装を最適化する impress top gearシリーズ

作者:Teiva Harsanyi,柴田芳樹
インプレス

Amazon

このパッケージの中には、io.Readerやio.Writerを受け取り処理する関数をテストする際に便利な関数なども用意されています。

今回は、iotest.TimeoutReader 関数について。

iotest.TimeoutReader 関数は、２回目のReadのときだけタイムアウトエラーを返す io.Reader です。通信処理のように複数回のReadが発生し、タイムアウトも考慮しなければならない処理のテストなどに使えます。個人的には、この iotest.TimeoutReader の存在を知れたのが一番嬉しい知識でした。これは便利。「２回目だけ」タイムアウトするって点も素晴らしい。

サンプル

random.go

package timeoutreader

import (
    "math/rand"
    "time"
)

const (
    charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
)

func randomString(buf []byte) {
    var (
        unixNano  = time.Now().UnixNano()
        rndSource = rand.NewSource(unixNano)
        rnd       = rand.New(rndSource)
    )

    for i := range buf {
        buf[i] = charset[rnd.Intn(len(charset))]
    }
}

read.go

package timeoutreader

import (
    "errors"
    "io"
    "time"
)

var (
    interval     = 10 * time.Millisecond
    ErrRetryOver = errors.New("retry over")
)

func readAllAtOnce(r io.Reader, p []byte) error {
    var (
        b   []byte
        err error
    )

    b, err = io.ReadAll(r)
    if err != nil {
        return err
    }

    copy(p, b)

    return nil
}

func readWithRetry(r io.Reader, p []byte, retries int) error {
    var (
        buf      = make([]byte, 1<<9)
        numRead  int
        offset   int
        count    int
        maxCount = retries + 1
        err      error
    )

    for count = 0; count < maxCount; {
        if len(p) <= offset {
            break
        }

        clear(buf)

        numRead, err = r.Read(buf)
        if err != nil {
            if errors.Is(err, io.EOF) {
                break
            }

            time.Sleep(interval)
            count++
            continue
        }

        copy(p[offset:offset+numRead], buf[:numRead])
        offset += numRead
    }

    if maxCount <= count {
        return ErrRetryOver
    }

    return nil
}

timeoutreader_test.go

package timeoutreader

import (
    "bytes"
    "io"
    "testing"
    "testing/iotest"
)

// TestReadAllAtOnce は、io.ReadAll()を利用してデータを読み取るテストケースです。
// テストで利用している io.Reader は、iotest.TimeoutReader()から生成しています。
//
// iotest.TimeoutReader() は、**２回目のみ**タイムアウトで失敗する
// io.Readerを作成してくれます。なので、通信処理などをテストする際にとても便利です。
// ２回目以降は普通にデータが読み取れるようになっています。
//
// このテストケースでは、読み取りに io.ReadAll() を利用しているため、２回目のエラーで
// 処理が返ってしまい、テストケースがFailとなります。io.ReadAll()には、リトライ処理などは
// 実装されていないため、タイムアウトが発生する可能性がある場面では利用しない方が良いということになります。
//
// # REFERENCES
//   - https://pkg.go.dev/testing/iotest@go1.23.0#TimeoutReader
func TestReadAllAtOnce(t *testing.T) {
    var (
        data   = make([]byte, 1<<10)
        buf    = make([]byte, len(data))
        reader io.Reader
        err    error
    )

    randomString(data)
    reader = iotest.TimeoutReader(bytes.NewReader(data))

    err = readAllAtOnce(reader, buf)
    if err != nil {
        t.Fatal(err)
    }

    if !bytes.Equal(data, buf) {
        t.Fatalf("[want] equal\t[got] not equal")
    }
}

// TestReadWithRetry は、リトライ処理を考慮した読み取り処理を利用してデータを読み取るテストケースです。
// テストで利用している io.Reader は、iotest.TimeoutReader()から生成しています。
//
// iotest.TimeoutReader() は、**２回目のみ**タイムアウトで失敗する
// io.Readerを作成してくれます。なので、通信処理などをテストする際にとても便利です。
// ２回目以降は普通にデータが読み取れるようになっています。
//
// このテストケースでは、読み取り中にタイムアウトを含むエラーが発生しても、指定回数リトライして
// 再試行するようになっているため、テストケースが通ります。
//
// # REFERENCES
//   - https://pkg.go.dev/testing/iotest@go1.23.0#TimeoutReader
func TestReadWithRetry(t *testing.T) {
    var (
        data   = make([]byte, 1<<10)
        buf    = make([]byte, len(data))
        reader io.Reader
        err    error
    )

    randomString(data)
    reader = iotest.TimeoutReader(bytes.NewReader(data))

    err = readWithRetry(reader, buf, 3)
    if err != nil {
        t.Fatal(err)
    }

    if !bytes.Equal(data, buf) {
        t.Fatalf("[want] equal\t[got] not equal")
    }
}

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

$ task
=== RUN   TestReadAllAtOnce
    timeoutreader_test.go:36: timeout
--- FAIL: TestReadAllAtOnce (0.00s)
=== RUN   TestReadWithRetry
--- PASS: TestReadWithRetry (0.01s)
FAIL
FAIL    github.com/devlights/try-golang/examples/singleapp/iotest/timeoutreader 0.013s
FAIL

io.ReadAll() を使っている方はエラーを検知してFailしていますが、リトライを考慮した方はSuccessしていますね。