You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
89 lines
1.6 KiB
Go
89 lines
1.6 KiB
Go
package limit
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/alicebob/miniredis"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/tal-tech/go-zero/core/logx"
|
|
"github.com/tal-tech/go-zero/core/stores/redis"
|
|
"github.com/tal-tech/go-zero/core/stores/redis/redistest"
|
|
)
|
|
|
|
func init() {
|
|
logx.Disable()
|
|
}
|
|
|
|
func TestTokenLimit_Rescue(t *testing.T) {
|
|
s, err := miniredis.Run()
|
|
assert.Nil(t, err)
|
|
|
|
const (
|
|
total = 100
|
|
rate = 5
|
|
burst = 10
|
|
)
|
|
l := NewTokenLimiter(rate, burst, redis.NewRedis(s.Addr(), redis.NodeType), "tokenlimit")
|
|
s.Close()
|
|
|
|
var allowed int
|
|
for i := 0; i < total; i++ {
|
|
time.Sleep(time.Second / time.Duration(total))
|
|
if i == total>>1 {
|
|
assert.Nil(t, s.Restart())
|
|
}
|
|
if l.Allow() {
|
|
allowed++
|
|
}
|
|
|
|
// make sure start monitor more than once doesn't matter
|
|
l.startMonitor()
|
|
}
|
|
|
|
assert.True(t, allowed >= burst+rate)
|
|
}
|
|
|
|
func TestTokenLimit_Take(t *testing.T) {
|
|
store, clean, err := redistest.CreateRedis()
|
|
assert.Nil(t, err)
|
|
defer clean()
|
|
|
|
const (
|
|
total = 100
|
|
rate = 5
|
|
burst = 10
|
|
)
|
|
l := NewTokenLimiter(rate, burst, store, "tokenlimit")
|
|
var allowed int
|
|
for i := 0; i < total; i++ {
|
|
time.Sleep(time.Second / time.Duration(total))
|
|
if l.Allow() {
|
|
allowed++
|
|
}
|
|
}
|
|
|
|
assert.True(t, allowed >= burst+rate)
|
|
}
|
|
|
|
func TestTokenLimit_TakeBurst(t *testing.T) {
|
|
store, clean, err := redistest.CreateRedis()
|
|
assert.Nil(t, err)
|
|
defer clean()
|
|
|
|
const (
|
|
total = 100
|
|
rate = 5
|
|
burst = 10
|
|
)
|
|
l := NewTokenLimiter(rate, burst, store, "tokenlimit")
|
|
var allowed int
|
|
for i := 0; i < total; i++ {
|
|
if l.Allow() {
|
|
allowed++
|
|
}
|
|
}
|
|
|
|
assert.True(t, allowed >= burst)
|
|
}
|