wangjia
3c687e5e0b
新包 internal/ratelimit：内存/redis(GCRA+Lua) 双实现 + 出错逐调用降级内存 （fail-open 到内存不 fail-closed）。REDIS_ADDR 空=内存模式，行为与既往一致； 配置后跨重启保状态、支持多实例。miniredis 全覆盖测试，零真实外部依赖。 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
102 lines
2.3 KiB
Go
102 lines
2.3 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// 迁自 middleware.TestKeyedLimiterSweepEvictsIdle：空闲 key 被 sweep 淘汰、活跃保留。
|
|
func TestMemLimiterSweepEvictsIdle(t *testing.T) {
|
|
l := newMemLimiter(rate.Limit(1), 1)
|
|
l.Allow("ip:a")
|
|
l.Allow("ip:b")
|
|
if len(l.entries) != 2 {
|
|
t.Fatalf("应有 2 个 entry，得到 %d", len(l.entries))
|
|
}
|
|
l.mu.Lock()
|
|
l.entries["ip:a"].lastSeen = time.Now().Add(-time.Hour)
|
|
l.mu.Unlock()
|
|
|
|
l.sweep(10 * time.Minute)
|
|
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
if _, ok := l.entries["ip:a"]; ok {
|
|
t.Fatal("空闲 key a 应被淘汰")
|
|
}
|
|
if _, ok := l.entries["ip:b"]; !ok {
|
|
t.Fatal("活跃 key b 不应被淘汰")
|
|
}
|
|
}
|
|
|
|
func TestMemLimiterBurst(t *testing.T) {
|
|
l := newMemLimiter(rate.Limit(1.0/60.0), 2) // 每分钟 1 个，突发 2
|
|
if !l.Allow("k").Allowed || !l.Allow("k").Allowed {
|
|
t.Fatal("burst 内应放行")
|
|
}
|
|
res := l.Allow("k")
|
|
if res.Allowed {
|
|
t.Fatal("超出 burst 应拒绝")
|
|
}
|
|
if res.RetryAfter <= 0 {
|
|
t.Fatal("拒绝时应给 RetryAfter")
|
|
}
|
|
if !l.Allow("other").Allowed {
|
|
t.Fatal("不同 key 独立桶")
|
|
}
|
|
}
|
|
|
|
func TestMemFailLockerLockAndExpire(t *testing.T) {
|
|
fl := newMemFailLocker()
|
|
const key = "S|u"
|
|
for i := 0; i < 3; i++ {
|
|
if fl.Locked(key) {
|
|
t.Fatalf("第 %d 次失败前不应锁定", i+1)
|
|
}
|
|
fl.RecordFailure(key, 3, 30*time.Millisecond)
|
|
}
|
|
if !fl.Locked(key) {
|
|
t.Fatal("达阈值应锁定")
|
|
}
|
|
time.Sleep(40 * time.Millisecond)
|
|
if fl.Locked(key) {
|
|
t.Fatal("锁定到期应自动解锁")
|
|
}
|
|
}
|
|
|
|
func TestMemFailLockerReset(t *testing.T) {
|
|
fl := newMemFailLocker()
|
|
fl.RecordFailure("k", 2, time.Minute)
|
|
fl.Reset("k")
|
|
fl.RecordFailure("k", 2, time.Minute)
|
|
if fl.Locked("k") {
|
|
t.Fatal("Reset 后重新计数，1 次失败不应锁定")
|
|
}
|
|
}
|
|
|
|
func TestMemFailLockerMaxZeroNeverLocks(t *testing.T) {
|
|
fl := newMemFailLocker()
|
|
for i := 0; i < 10; i++ {
|
|
fl.RecordFailure("k", 0, time.Minute)
|
|
}
|
|
if fl.Locked("k") {
|
|
t.Fatal("max<=0 不应锁定")
|
|
}
|
|
}
|
|
|
|
func TestMemCounterIncrAndTTL(t *testing.T) {
|
|
c := newMemCounter()
|
|
n1, _ := c.Incr("k", 30*time.Millisecond)
|
|
n2, _ := c.Incr("k", 30*time.Millisecond)
|
|
if n1 != 1 || n2 != 2 {
|
|
t.Fatalf("累加应为 1,2，得到 %d,%d", n1, n2)
|
|
}
|
|
time.Sleep(40 * time.Millisecond)
|
|
n3, _ := c.Incr("k", 30*time.Millisecond)
|
|
if n3 != 1 {
|
|
t.Fatalf("TTL 过期后应重新从 1 计，得到 %d", n3)
|
|
}
|
|
}
|