package ratelimit import "testing" func TestAllowRequestDaily(t *testing.T) { l := New() if !l.AllowRequestDaily(1, 2) { t.Fatal("first request should be allowed") } if !l.AllowRequestDaily(1, 2) { t.Fatal("second request should be allowed") } if l.AllowRequestDaily(1, 2) { t.Fatal("third request should be blocked") } // 另一个 key 不受影响 if !l.AllowRequestDaily(2, 2) { t.Fatal("other key should be allowed") } } func TestTokensDaily(t *testing.T) { l := New() l.AddTokens(1, 100) l.AddTokens(1, 50) if got := l.TokensUsed(1); got != 150 { t.Fatalf("tokens = %d, want 150", got) } if got := l.TokensUsed(2); got != 0 { t.Fatalf("other key tokens = %d, want 0", got) } // 配额检查语义:已达 150 再设配额 150 应拒绝 l.AddTokens(3, 150) if l.TokensUsed(3) >= 150 { // 注意:>= 表示已达上限,Auth 层据此拒绝 } } func TestAllowUserRate(t *testing.T) { l := New() // limit=1:立即第二次应被拒绝 if !l.AllowUserRate(9, 1) { t.Fatal("first should be allowed") } if l.AllowUserRate(9, 1) { t.Fatal("second in same second should be blocked") } // limit<=0 不限制 if !l.AllowUserRate(9, 0) { t.Fatal("limit 0 should always allow") } }