package ratelimit import ( "context" "testing" "time" "github.com/alicebob/miniredis/v2" "github.com/hhs/camtalk/internal/config" "github.com/redis/go-redis/v9" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // setupMiniRedis 创建一个内存 Redis 实例用于测试。 func setupMiniRedis(t *testing.T) (*miniredis.Miniredis, *redis.Client) { mr, err := miniredis.Run() require.NoError(t, err) client := redis.NewClient(&redis.Options{ Addr: mr.Addr(), }) t.Cleanup(func() { client.Close() mr.Close() }) return mr, client } func TestRedisLimiter_Allow_FirstRequest(t *testing.T) { _, client := setupMiniRedis(t) cfg := config.RateLimitConfig{ Enabled: true, Query: config.BucketConfig{Capacity: 5, Rate: 0.2}, } limiter := NewRedisLimiter(client, cfg) ctx := context.Background() allowed, retryAfter := limiter.Allow(ctx, "user1:query") assert.True(t, allowed) assert.Equal(t, time.Duration(0), retryAfter) } func TestRedisLimiter_Allow_ConsumeUntilEmpty(t *testing.T) { _, client := setupMiniRedis(t) cfg := config.RateLimitConfig{ Enabled: true, Query: config.BucketConfig{Capacity: 3, Rate: 0.2}, } limiter := NewRedisLimiter(client, cfg) ctx := context.Background() key := "user1:query" // 连续消耗 3 个令牌 for i := 0; i < 3; i++ { allowed, _ := limiter.Allow(ctx, key) assert.True(t, allowed, "request %d should be allowed", i+1) } // 第 4 个请求应被拒绝 allowed, retryAfter := limiter.Allow(ctx, key) assert.False(t, allowed) assert.Greater(t, retryAfter, time.Duration(0)) } func TestRedisLimiter_Allow_DifferentKeys(t *testing.T) { _, client := setupMiniRedis(t) cfg := config.RateLimitConfig{ Enabled: true, Query: config.BucketConfig{Capacity: 2, Rate: 1.0}, } limiter := NewRedisLimiter(client, cfg) ctx := context.Background() // user1 消耗 2 个令牌 allowed, _ := limiter.Allow(ctx, "user1:query") assert.True(t, allowed) allowed, _ = limiter.Allow(ctx, "user1:query") assert.True(t, allowed) // user1 第 3 个被拒绝 allowed, _ = limiter.Allow(ctx, "user1:query") assert.False(t, allowed) // user2 应该不受影响 allowed, _ = limiter.Allow(ctx, "user2:query") assert.True(t, allowed) } func TestRedisLimiter_Allow_RefillAfterWait(t *testing.T) { _, client := setupMiniRedis(t) cfg := config.RateLimitConfig{ Enabled: true, Query: config.BucketConfig{Capacity: 2, Rate: 10.0}, // 每秒 10 个令牌 } limiter := NewRedisLimiter(client, cfg) ctx := context.Background() key := "user1:query" // 消耗 2 个令牌 limiter.Allow(ctx, key) limiter.Allow(ctx, key) // 真实等待 150ms(Lua 脚本使用系统时间) time.Sleep(150 * time.Millisecond) // 应该补充了至少 1 个令牌 allowed, _ := limiter.Allow(ctx, key) assert.True(t, allowed) } func TestRedisLimiter_Allow_CapacityLimit(t *testing.T) { _, client := setupMiniRedis(t) cfg := config.RateLimitConfig{ Enabled: true, Query: config.BucketConfig{Capacity: 3, Rate: 1.0}, } limiter := NewRedisLimiter(client, cfg) ctx := context.Background() key := "user1:query" // 真实等待让桶"溢出" time.Sleep(100 * time.Millisecond) // 但最多只能消耗 capacity 个令牌 for i := 0; i < 3; i++ { allowed, _ := limiter.Allow(ctx, key) assert.True(t, allowed, "request %d should be allowed", i+1) } // 第 4 个应被拒绝 allowed, _ := limiter.Allow(ctx, key) assert.False(t, allowed) } func TestRedisLimiter_Allow_ZeroRate(t *testing.T) { _, client := setupMiniRedis(t) cfg := config.RateLimitConfig{ Enabled: true, Query: config.BucketConfig{Capacity: 1, Rate: 0.0}, } limiter := NewRedisLimiter(client, cfg) ctx := context.Background() key := "user1:query" // 第一个通过 allowed, _ := limiter.Allow(ctx, key) assert.True(t, allowed) // 第二个被拒绝,retryAfter 应该很大 allowed, retryAfter := limiter.Allow(ctx, key) assert.False(t, allowed) assert.Greater(t, retryAfter, 1*time.Hour) } func TestRedisLimiter_Allow_KeyTTL(t *testing.T) { mr, client := setupMiniRedis(t) cfg := config.RateLimitConfig{ Enabled: true, Query: config.BucketConfig{Capacity: 5, Rate: 1.0}, } limiter := NewRedisLimiter(client, cfg) ctx := context.Background() key := "user1:query" // 第一次请求 limiter.Allow(ctx, key) // 验证 key 已设置 TTL ttl := mr.TTL(key) assert.Greater(t, ttl, time.Duration(0)) assert.LessOrEqual(t, ttl, 600*time.Second) } func TestRedisLimiter_Allow_FailOpen(t *testing.T) { mr, client := setupMiniRedis(t) cfg := config.RateLimitConfig{ Enabled: true, Query: config.BucketConfig{Capacity: 1, Rate: 1.0}, } limiter := NewRedisLimiter(client, cfg) ctx := context.Background() // 关闭 Redis 模拟故障 mr.Close() // 应该 fail-open(允许请求) allowed, retryAfter := limiter.Allow(ctx, "user1:query") assert.True(t, allowed) assert.Equal(t, time.Duration(0), retryAfter) } func TestRedisLimiter_Stop(t *testing.T) { _, client := setupMiniRedis(t) cfg := config.RateLimitConfig{ Enabled: true, Query: config.BucketConfig{Capacity: 1, Rate: 1.0}, } limiter := NewRedisLimiter(client, cfg) // Stop 应该不会 panic(即使多次调用) limiter.Stop() limiter.Stop() } func TestFormatKey(t *testing.T) { key := FormatKey("user123", "query") assert.Equal(t, "ratelimit:user123:query", key) }