128 lines
3.0 KiB
Go
128 lines
3.0 KiB
Go
package repo
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/agentmail/gateway/internal/db"
|
||
)
|
||
|
||
// 新建会话速率限制：DB 版
|
||
//
|
||
// 这些测试用真实的 SQLite（setupTestDB），验证速率限制的原子性与窗口滑动。
|
||
// 原来的内存版测试依赖 sessionRateLimiter 结构体，替换为 DB 版后重写。
|
||
|
||
func TestSessionRateAllowsUpToLimit(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
for i := 1; i <= sessionRateLimit; i++ {
|
||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||
t.Fatalf("第 %d 次应放行（上限 %d）", i, sessionRateLimit)
|
||
}
|
||
}
|
||
ok, retry := AllowNewSession(ctx, "bot")
|
||
if ok {
|
||
t.Fatal("超过上限应拦下")
|
||
}
|
||
if retry < 1 {
|
||
t.Fatalf("应给出正的重试等待秒数，实际 %d", retry)
|
||
}
|
||
}
|
||
|
||
func TestSessionRateIsPerAgent(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
for i := 0; i < sessionRateLimit; i++ {
|
||
AllowNewSession(ctx, "busy")
|
||
}
|
||
if ok, _ := AllowNewSession(ctx, "busy"); ok {
|
||
t.Fatal("busy 应已被拦")
|
||
}
|
||
if ok, _ := AllowNewSession(ctx, "idle"); !ok {
|
||
t.Fatal("另一个 Agent 不该被牵连")
|
||
}
|
||
}
|
||
|
||
// 并发请求不能把上限刷穿（判断与记账必须原子）
|
||
func TestSessionRateConcurrentDoesNotOverrun(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
var wg sync.WaitGroup
|
||
var mu sync.Mutex
|
||
passed := 0
|
||
|
||
for i := 0; i < sessionRateLimit*4; i++ {
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
if ok, _ := AllowNewSession(ctx, "bot"); ok {
|
||
mu.Lock()
|
||
passed++
|
||
mu.Unlock()
|
||
}
|
||
}()
|
||
}
|
||
wg.Wait()
|
||
|
||
if passed != sessionRateLimit {
|
||
t.Fatalf("%d 并发下放行 %d 次，期望恰好 %d 次",
|
||
sessionRateLimit*4, passed, sessionRateLimit)
|
||
}
|
||
}
|
||
|
||
// 建会话失败时要还名额
|
||
func TestSessionRateRelease(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
for i := 0; i < sessionRateLimit; i++ {
|
||
AllowNewSession(ctx, "bot")
|
||
}
|
||
if ok, _ := AllowNewSession(ctx, "bot"); ok {
|
||
t.Fatal("应已刷满")
|
||
}
|
||
ReleaseNewSession(ctx, "bot")
|
||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||
t.Fatal("归还名额后应能再开一条")
|
||
}
|
||
}
|
||
|
||
// 人类不走限速（空 agentName）
|
||
func TestAllowNewSessionSkipsHumans(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
for i := 0; i < sessionRateLimit*3; i++ {
|
||
if ok, _ := AllowNewSession(ctx, ""); !ok {
|
||
t.Fatal("人类不该被限速")
|
||
}
|
||
}
|
||
}
|
||
|
||
// 窗口滑过后自动恢复（过期记录自动清理）
|
||
func TestSessionRateWindowSlides(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
// 插入过期记录（1小时前）
|
||
cutoff := time.Now().Add(-sessionRateWindow - time.Minute)
|
||
for i := 0; i < sessionRateLimit; i++ {
|
||
_, err := db.DB.ExecContext(ctx,
|
||
`INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`,
|
||
"session:bot", cutoff)
|
||
if err != nil {
|
||
t.Fatalf("插入过期记录失败: %v", err)
|
||
}
|
||
}
|
||
|
||
// 窗口外的记录应被清理，此次应放行
|
||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||
t.Fatal("窗口外的记录应被清掉，此次应放行")
|
||
}
|
||
}
|