36 lines
713 B
Go
36 lines
713 B
Go
package middleware
|
||
|
||
import (
|
||
"testing"
|
||
"time"
|
||
)
|
||
|
||
func TestRateLimiter_AllowAndRefill(t *testing.T) {
|
||
rl := NewRateLimiter(2, 1) // 容量2，每秒补充1
|
||
|
||
// 注入可控时间
|
||
now := time.Unix(0, 0).UnixNano()
|
||
rl.now = func() int64 { return now }
|
||
rl.lastRefill = now
|
||
rl.tokens = 2
|
||
|
||
if !rl.Allow() {
|
||
t.Fatalf("expected first Allow() true")
|
||
}
|
||
if !rl.Allow() {
|
||
t.Fatalf("expected second Allow() true")
|
||
}
|
||
if rl.Allow() {
|
||
t.Fatalf("expected third Allow() false (no tokens)")
|
||
}
|
||
|
||
// 过 1 秒应补充 1 个令牌
|
||
now += int64(time.Second)
|
||
if !rl.Allow() {
|
||
t.Fatalf("expected Allow() true after refill")
|
||
}
|
||
if rl.Allow() {
|
||
t.Fatalf("expected Allow() false again (tokens should be 0)")
|
||
}
|
||
}
|