454 lines
13 KiB
Go
454 lines
13 KiB
Go
package autoassignment
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/alicebob/miniredis/v2"
|
|
"github.com/redis/go-redis/v9"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// setupTestRedis 创建一个基于 miniredis 的测试 Redis 环境。
|
|
// 返回 miniredis 服务实例和 redis.Client。
|
|
// miniredis.RunT 会自动在 t.Cleanup 中关闭服务。
|
|
func setupTestRedis(t *testing.T) (*miniredis.Miniredis, *redis.Client) {
|
|
t.Helper()
|
|
|
|
mr := miniredis.RunT(t)
|
|
rdb := redis.NewClient(&redis.Options{
|
|
Addr: mr.Addr(),
|
|
})
|
|
|
|
t.Cleanup(func() {
|
|
rdb.Close()
|
|
})
|
|
|
|
return mr, rdb
|
|
}
|
|
|
|
// --- IsAllowed 测试 ---
|
|
|
|
// TestIsAllowed_首次请求允许 验证：新窗口中没有任何计数时，IsAllowed 应返回 true。
|
|
func TestIsAllowed_首次请求允许(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
limit := 3
|
|
windowSeconds := 300
|
|
|
|
// 新窗口中没有任何计数，应该允许
|
|
assert.True(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
|
"首次请求应被允许")
|
|
}
|
|
|
|
// TestIsAllowed_达到限制后拒绝 验证：当计数达到 limit 后，IsAllowed 应返回 false。
|
|
func TestIsAllowed_达到限制后拒绝(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
limit := 3
|
|
windowSeconds := 300
|
|
|
|
// 先 Increment 3 次，达到限制
|
|
for i := 0; i < limit; i++ {
|
|
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// 达到限制后，应该拒绝
|
|
assert.False(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
|
"达到限制后应被拒绝")
|
|
}
|
|
|
|
// TestIsAllowed_未达到限制时允许 验证：计数小于 limit 时，IsAllowed 应返回 true。
|
|
func TestIsAllowed_未达到限制时允许(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
limit := 3
|
|
windowSeconds := 300
|
|
|
|
// Increment 2 次，未达到限制
|
|
for i := 0; i < limit-1; i++ {
|
|
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// 未达到限制，应该允许
|
|
assert.True(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
|
"未达到限制时应被允许")
|
|
}
|
|
|
|
// TestIsAllowed_Limit为零时总是允许 验证：limit <= 0 时，IsAllowed 应返回 true（不限制）。
|
|
func TestIsAllowed_Limit为零时总是允许(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
|
|
// limit = 0，应总是允许
|
|
assert.True(t, rl.IsAllowed(ctx, 1, 10, 0, 300),
|
|
"limit 为 0 时应总是允许")
|
|
|
|
// limit = -1，应总是允许
|
|
assert.True(t, rl.IsAllowed(ctx, 1, 10, -1, 300),
|
|
"limit 为负数时应总是允许")
|
|
}
|
|
|
|
// TestIsAllowed_不同Agent独立计数 验证：不同 agent 的限制独立计算。
|
|
func TestIsAllowed_不同Agent独立计数(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
limit := 2
|
|
windowSeconds := 300
|
|
|
|
// Agent A 达到限制
|
|
for i := 0; i < limit; i++ {
|
|
err := rl.Increment(ctx, inboxID, 10, windowSeconds)
|
|
require.NoError(t, err)
|
|
}
|
|
assert.False(t, rl.IsAllowed(ctx, inboxID, 10, limit, windowSeconds),
|
|
"Agent A 达到限制后应被拒绝")
|
|
|
|
// Agent B 未达到限制，应允许
|
|
assert.True(t, rl.IsAllowed(ctx, inboxID, 20, limit, windowSeconds),
|
|
"Agent B 应独立计算，首次请求应被允许")
|
|
}
|
|
|
|
// TestIsAllowed_不同Inbox独立计数 验证：不同 inbox 的限制独立计算。
|
|
func TestIsAllowed_不同Inbox独立计数(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
agentID := uint(10)
|
|
limit := 2
|
|
windowSeconds := 300
|
|
|
|
// Inbox 1 达到限制
|
|
for i := 0; i < limit; i++ {
|
|
err := rl.Increment(ctx, 1, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
}
|
|
assert.False(t, rl.IsAllowed(ctx, 1, agentID, limit, windowSeconds),
|
|
"Inbox 1 达到限制后应被拒绝")
|
|
|
|
// Inbox 2 未达到限制，应允许
|
|
assert.True(t, rl.IsAllowed(ctx, 2, agentID, limit, windowSeconds),
|
|
"Inbox 2 应独立计算，首次请求应被允许")
|
|
}
|
|
|
|
// --- Increment 测试 ---
|
|
|
|
// TestIncrement_正常计数增加 验证：每次 Increment 后计数应递增 1。
|
|
func TestIncrement_正常计数增加(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
windowSeconds := 300
|
|
|
|
// 验证初始计数为 0
|
|
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 0, count)
|
|
|
|
// Increment 1 次
|
|
err = rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
|
|
count, err = rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 1, count)
|
|
|
|
// Increment 第 2 次
|
|
err = rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
|
|
count, err = rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 2, count)
|
|
}
|
|
|
|
// TestIncrement_设置TTL 验证：Increment 后 key 应有正确的 TTL。
|
|
func TestIncrement_设置TTL(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
windowSeconds := 300
|
|
|
|
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
|
|
// 验证 key 存在且有 TTL
|
|
key := rl.key(inboxID, agentID, windowSeconds)
|
|
ttl, err := rdb.TTL(ctx, key).Result()
|
|
require.NoError(t, err)
|
|
|
|
// TTL 应大于 0，且不超过 windowSeconds
|
|
assert.Greater(t, ttl, time.Duration(0))
|
|
assert.LessOrEqual(t, ttl, time.Duration(windowSeconds)*time.Second)
|
|
}
|
|
|
|
// TestIncrement_重复Increment不覆盖TTL 验证：窗口内后续 Increment 不应重新设置 TTL。
|
|
func TestIncrement_重复Increment不覆盖TTL(t *testing.T) {
|
|
mr, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
windowSeconds := 60
|
|
|
|
// 第一次 Increment，设置 TTL
|
|
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
|
|
key := rl.key(inboxID, agentID, windowSeconds)
|
|
ttl1, err := rdb.TTL(ctx, key).Result()
|
|
require.NoError(t, err)
|
|
|
|
// 快进时间，使 TTL 减少
|
|
mr.FastForward(30 * time.Second)
|
|
|
|
ttlAfterForward, err := rdb.TTL(ctx, key).Result()
|
|
require.NoError(t, err)
|
|
assert.Less(t, ttlAfterForward, ttl1, "快进后 TTL 应减少")
|
|
|
|
// 第二次 Increment，不应重新设置 TTL（TTL 应继续减少而非恢复）
|
|
err = rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
|
|
ttl2, err := rdb.TTL(ctx, key).Result()
|
|
require.NoError(t, err)
|
|
assert.Less(t, ttl2, ttl1, "后续 Increment 不应重新设置 TTL")
|
|
}
|
|
|
|
// --- Redis 错误降级测试 ---
|
|
|
|
// TestIsAllowed_Redis错误时允许 验证：当 Redis 出错时，IsAllowed 应返回 true（降级允许）。
|
|
func TestIsAllowed_Redis错误时允许(t *testing.T) {
|
|
mr, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
limit := 1
|
|
windowSeconds := 300
|
|
|
|
// 先 Increment 1 次，达到限制
|
|
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.False(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
|
"正常情况下达到限制应拒绝")
|
|
|
|
// 关闭 miniredis，模拟 Redis 连接错误
|
|
mr.Close()
|
|
|
|
// Redis 错误时应降级允许
|
|
assert.True(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
|
"Redis 错误时应降级允许请求")
|
|
}
|
|
|
|
// TestIncrement_Redis错误返回错误 验证：当 Redis 出错时，Increment 应返回错误。
|
|
func TestIncrement_Redis错误返回错误(t *testing.T) {
|
|
mr, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
windowSeconds := 300
|
|
|
|
// 关闭 miniredis，模拟 Redis 连接错误
|
|
mr.Close()
|
|
|
|
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
assert.Error(t, err, "Redis 错误时 Increment 应返回错误")
|
|
}
|
|
|
|
// --- 窗口过期后限制重置测试 ---
|
|
|
|
// TestIsAllowed_窗口过期后限制重置 验证：时间窗口过期后，计数应重置，请求应重新被允许。
|
|
func TestIsAllowed_窗口过期后限制重置(t *testing.T) {
|
|
mr, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
limit := 2
|
|
windowSeconds := 60
|
|
|
|
// Increment 达到限制
|
|
for i := 0; i < limit; i++ {
|
|
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
}
|
|
assert.False(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
|
"达到限制后应被拒绝")
|
|
|
|
// 快进时间超过窗口时长，使 key 过期
|
|
mr.FastForward(time.Duration(windowSeconds) * time.Second)
|
|
|
|
// 窗口过期后，应重新允许
|
|
assert.True(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
|
"窗口过期后应重新允许请求")
|
|
}
|
|
|
|
// --- GetCount 测试 ---
|
|
|
|
// TestGetCount_无计数时返回零 验证：没有 Increment 时 GetCount 应返回 0。
|
|
func TestGetCount_无计数时返回零(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
windowSeconds := 300
|
|
|
|
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 0, count, "无计数时应返回 0")
|
|
}
|
|
|
|
// TestGetCount_正确返回计数 验证：Increment 后 GetCount 应返回正确的计数。
|
|
func TestGetCount_正确返回计数(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
agentID := uint(10)
|
|
windowSeconds := 300
|
|
|
|
for i := 1; i <= 5; i++ {
|
|
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
|
|
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, i, count)
|
|
}
|
|
}
|
|
|
|
// --- Reset 测试 ---
|
|
|
|
// TestReset_清除所有计数 验证：Reset 应清除指定 inbox 下所有 agent 的计数。
|
|
func TestReset_清除所有计数(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
inboxID := uint(1)
|
|
windowSeconds := 300
|
|
|
|
// 为多个 agent Increment
|
|
for agentID := uint(10); agentID <= 13; agentID++ {
|
|
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// 验证计数存在
|
|
for agentID := uint(10); agentID <= 13; agentID++ {
|
|
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 1, count)
|
|
}
|
|
|
|
// Reset inbox
|
|
err := rl.Reset(ctx, inboxID)
|
|
require.NoError(t, err)
|
|
|
|
// 验证计数已被清除
|
|
for agentID := uint(10); agentID <= 13; agentID++ {
|
|
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 0, count, "Reset 后计数应被清除")
|
|
}
|
|
}
|
|
|
|
// TestReset_不影响其他Inbox 验证：Reset 一个 inbox 不应影响其他 inbox 的计数。
|
|
func TestReset_不影响其他Inbox(t *testing.T) {
|
|
_, rdb := setupTestRedis(t)
|
|
rl := NewRateLimiter(rdb)
|
|
|
|
ctx := context.Background()
|
|
windowSeconds := 300
|
|
|
|
// Inbox 1 和 Inbox 2 都有计数
|
|
err := rl.Increment(ctx, 1, 10, windowSeconds)
|
|
require.NoError(t, err)
|
|
err = rl.Increment(ctx, 2, 10, windowSeconds)
|
|
require.NoError(t, err)
|
|
|
|
// Reset Inbox 1
|
|
err = rl.Reset(ctx, 1)
|
|
require.NoError(t, err)
|
|
|
|
// Inbox 1 计数被清除
|
|
count, err := rl.GetCount(ctx, 1, 10, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 0, count)
|
|
|
|
// Inbox 2 计数不受影响
|
|
count, err = rl.GetCount(ctx, 2, 10, windowSeconds)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 1, count, "其他 inbox 的计数不应受影响")
|
|
}
|
|
|
|
// --- ParseAgentID 测试 ---
|
|
|
|
// TestParseAgentID_正常解析 验证：ParseAgentID 应正确解析合法字符串。
|
|
func TestParseAgentID_正常解析(t *testing.T) {
|
|
id, err := ParseAgentID("42")
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, uint(42), id)
|
|
}
|
|
|
|
// TestParseAgentID_非法字符串 验证：ParseAgentID 应对非法字符串返回错误。
|
|
func TestParseAgentID_非法字符串(t *testing.T) {
|
|
_, err := ParseAgentID("abc")
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// TestParseAgentID_空字符串 验证：ParseAgentID 应对空字符串返回错误。
|
|
func TestParseAgentID_空字符串(t *testing.T) {
|
|
_, err := ParseAgentID("")
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// TestParseAgentID_负数 验证：ParseAgentID 应对负数字符串返回错误。
|
|
func TestParseAgentID_负数(t *testing.T) {
|
|
_, err := ParseAgentID("-1")
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// TestParseAgentID_零 验证：ParseAgentID 应正确解析 "0"。
|
|
func TestParseAgentID_零(t *testing.T) {
|
|
id, err := ParseAgentID("0")
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, uint(0), id)
|
|
} |