openbilibili/library/rate/limit/limit_test.go

package limit

import (
	"context"
	"testing"
	"time"

	"go-common/library/rate"
)

func worker(qps int64, ch chan struct{}) {
	for {
		<-ch
		time.Sleep(time.Duration(int64(time.Second) / qps))
	}
}

func TestRateSuccess(t *testing.T) {
	ch := make(chan struct{})
	go worker(100, ch)
	failed := producer(New(nil), 100, ch)
	if failed > 0 {
		t.Fatalf("Should be rejected 0 time,but (%d)", failed)
	}
}

func TestRateFail(t *testing.T) {
	ch := make(chan struct{})
	go worker(100, ch)
	failed := producer(New(nil), 200, ch)
	if failed < 900 {
		t.Fatalf("Should be rejected more than 900 times,but (%d)", failed)
	}
}

func TestRateFailMuch(t *testing.T) {
	ch := make(chan struct{})
	go worker(10, ch)
	failed := producer(New(nil), 200, ch)
	if failed < 1600 {
		t.Fatalf("Should be rejected more than 1600 times,but (%d)", failed)
	}
}

func producer(l *Limiter, qps int64, ch chan struct{}) (failed int) {
	for i := 0; i < int(qps)*10; i++ {
		go func() {
			done, err := l.Allow(context.Background())
			defer done(rate.Success)
			if err == nil {
				ch <- struct{}{}
			} else {
				failed++
			}
		}()
		time.Sleep(time.Duration(int64(time.Second) / qps))
	}
	return
}