CamTalk/backend/internal/ratelimit/redis_bucket_test.go

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)
}