weichen
29f14ac5f2
Details (original message): Add multi-instance groundwork: blobstore seam for artifacts, PG LISTEN/NOTIFY event fanout, and per-tenant rate limiting with native Redis backend
58 lines
1.4 KiB
Go
58 lines
1.4 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestMemoryWindow(t *testing.T) {
|
|
m := NewMemory()
|
|
ctx := context.Background()
|
|
for i := 1; i <= 3; i++ {
|
|
if !m.Allow(ctx, "o", 3, time.Minute) {
|
|
t.Fatalf("request %d should pass", i)
|
|
}
|
|
}
|
|
if m.Allow(ctx, "o", 3, time.Minute) {
|
|
t.Fatal("4th request must be limited")
|
|
}
|
|
// 其他租户互不影响。
|
|
if !m.Allow(ctx, "p", 3, time.Minute) {
|
|
t.Fatal("other tenant must pass")
|
|
}
|
|
// 窗口滚动后恢复。
|
|
m.now = func() time.Time { return time.Now().Add(2 * time.Minute) }
|
|
if !m.Allow(ctx, "o", 3, time.Minute) {
|
|
t.Fatal("new window must reset")
|
|
}
|
|
// limit<=0 = 不限。
|
|
if !m.Allow(ctx, "x", 0, time.Minute) {
|
|
t.Fatal("limit 0 means unlimited")
|
|
}
|
|
}
|
|
|
|
// Redis 共享档（环境门控：AGENT_REDIS_TEST_ADDR，本机原生实例）。
|
|
func TestRedisWindow(t *testing.T) {
|
|
addr := os.Getenv("AGENT_REDIS_TEST_ADDR")
|
|
if addr == "" {
|
|
t.Skip("AGENT_REDIS_TEST_ADDR not set")
|
|
}
|
|
r := NewRedis(addr, "", 15) // 独立 db 避免污染
|
|
defer r.Close()
|
|
if err := r.Ping(context.Background()); err != nil {
|
|
t.Fatalf("ping: %v", err)
|
|
}
|
|
ctx := context.Background()
|
|
key := "test-" + time.Now().Format("150405.000000000")
|
|
for i := 1; i <= 2; i++ {
|
|
if !r.Allow(ctx, key, 2, time.Minute) {
|
|
t.Fatalf("request %d should pass", i)
|
|
}
|
|
}
|
|
if r.Allow(ctx, key, 2, time.Minute) {
|
|
t.Fatal("3rd request must be limited")
|
|
}
|
|
}
|