60 lines
1.1 KiB
Go
60 lines
1.1 KiB
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
|
|
}
|