@kyanny's blog

創造性は記憶によるところが大きい - 黒澤明

Go: golang.org/x/time/rate を利用したコードのテストの書き方

速度制限をかけているコードの実行に要した処理時間を計測すれば良い。

package main

import (
    "context"
    "fmt"
    "testing"
    "time"

    "golang.org/x/time/rate"
)

func TestWait(t *testing.T) {
    // rate = 10 req / sec
    // burst = 1
    // 0.1 秒ごとに token が 1 ずつ増える
    r := rate.Every(time.Second / 10)
    b := 1
    limiter := rate.NewLimiter(r, b)

    ctx := context.Background()

    tests := []struct {
        name    string
        n int
    }{
        {
            name: "n = 10",
            n: 10,
        },
        // ループの実行開始直後に token を 1 消費する (burst=1)
        // これを考慮に入れるため、 PASS するテストケースではループを 11 回実行している (n=11)
        {
            name: "n = 11",
            n: 11,
        },
    }
    for _, tt := range tests {
        t.Run(tt.name, func(t *testing.T) {
            t0 := time.Now()
            for i := 0; i < tt.n; i++ {
                // token が増えるとすぐにループが実行され、 token を使い切る
                // token を使い切っていると limiter.Wait でブロックする
                // つまり、ループを 1 回実行するのに 0.1 秒かかる
                if err := limiter.Wait(ctx); err != nil {
                    // noop
                }
                // ここ (Wait() のあと) に速度制限をかけたい処理を書く
                // 例: resp, err := http.DefaultClient.Get("http://example.com/")
            }
            t1 := time.Now()
            elapsed := t1.Sub(t0).Seconds()
            // ループを 10 回実行する間に limiter.Wait でブロックする時間の合計は (10-1) * 0.1 = 0.9 秒
            // ループを 11 回実行する間に limiter.Wait でブロックする時間の合計は (11-1) * 0.1 = 1.0 秒
            // ループを 11 回実行するのに要した時間 (elapsed) が 1.0 秒以上であれば 10 req / sec の速度制限が効いているといえる
            if elapsed < 1.0 {
                t.Errorf("must take > 1.0 sec, but took %f sec", elapsed)
            }
            fmt.Printf("took %f sec\n", elapsed)
        })
    }
}

https://play.golang.org/p/lnrqdM3FI5c