package repo import ( "context" "testing" "time" "github.com/agentmail/gateway/internal/db" ) /* 429 的 retry_after —— 2026-09-28。 RateLimitCheckAndRecord 撞上限时算「最早那条 + window = 什么时候能再试」。 它原先用 `SELECT MIN(ts)`,SQLite 驱动把聚合结果当 string 返回, 扫进 *time.Time 失败 → 落到兜底 `return false, 60`。 后果是所有 429 的 retry_after 恒为 60 秒:新建会话与日历事件的窗口都是 1 小时,报错却说「60 秒后再试」。**限流本身是好的**(该拒还是拒), 坏的是它给出的数字 —— 压测/客户端拿它算退避节奏会算出一张假的时间表。 判据刻意不写「retry_after > 60」这种弱断言(60 也会碰巧通过一部分), 而是钉住语义:**retry_after 必须接近真实剩余窗口**,且不得超过它。 */ // TestRateLimitRetryAfterReflectsRealWindow 是本文件的主判据。 // 修复前:retry=60(与 window 无关);修复后:≈window。 func TestRateLimitRetryAfterReflectsRealWindow(t *testing.T) { setupTestDB(t) ctx := context.Background() const window = time.Hour const limit = 3 // 桶里已有 limit-1 条、且最早的一条在 10 分钟前 ⇒ 真实剩余约 50 分钟 seedRateLimits(t, "probe", limit-1, 10*time.Minute) // 这一条应当放行(还没到上限) if ok, retry := RateLimitCheckAndRecord(ctx, "probe", window, limit); !ok { t.Fatalf("第 %d 条应放行,却拿到 retry=%d", limit, retry) } // 第 limit+1 条必须被拒,且 retry_after 接近真实剩余(≈50min),绝不是 60 ok, retry := RateLimitCheckAndRecord(ctx, "probe", window, limit) if ok { t.Fatal("超过上限仍放行") } if retry == 60 { t.Errorf("retry_after=60(MIN(ts) 扫描失败的兜底值),期望≈%d", int((window - 10*time.Minute).Seconds())) } wantApprox := int((window - 10*time.Minute).Seconds()) if rlAbs(retry-wantApprox) > 5 { t.Errorf("retry_after=%d,与真实剩余 %d 差太多", retry, wantApprox) } // 上界也不能越界:报的时间超过真实剩余 = 让调用方空等 if retry > wantApprox+2 { t.Errorf("retry_after=%d 超过真实剩余 %d,会让调用方多等", retry, wantApprox) } } // TestRateLimitRetryAfterNeverExceedsWindow 钉住上界: // 桶被塞满的瞬间就该是 retry≈window。 func TestRateLimitRetryAfterNeverExceedsWindow(t *testing.T) { setupTestDB(t) ctx := context.Background() const window = time.Hour const limit = 2 seedRateLimits(t, "probe", limit, 0) // 最早的一条就在此刻 ok, retry := RateLimitCheckAndRecord(ctx, "probe", window, limit) if ok { t.Fatal("超过上限仍放行") } if retry > int(window.Seconds())+2 { t.Errorf("retry_after=%d 超过窗口 %ds", retry, int(window.Seconds())) } if retry < int(window.Seconds())-5 { t.Errorf("retry_after=%d 明显小于窗口 %ds(真值应≈窗口)", retry, int(window.Seconds())) } } // TestRateLimitAllowsAfterWindowRolls 确认窗口滚动后自动放行 —— // 429 不是「永久封禁」。 func TestRateLimitAllowsAfterWindowRolls(t *testing.T) { setupTestDB(t) ctx := context.Background() const window = time.Minute seedRateLimits(t, "probe", 5, 2*time.Minute) // 全部在窗口外 if ok, retry := RateLimitCheckAndRecord(ctx, "probe", window, 5); !ok { t.Fatalf("窗口已滚过却仍被拒,retry=%d", retry) } } // TestRateLimitCountsOnlyRecordsInsideWindow 是「窗口外不该计数」的方向。 // 少了这条,一次陈旧爆发就能永久压住一个桶。 func TestRateLimitCountsOnlyRecordsInsideWindow(t *testing.T) { setupTestDB(t) ctx := context.Background() const window = time.Minute // 100 条老记录 + 窗口内 0 条 seedRateLimits(t, "probe", 100, 5*time.Minute) for i := 0; i < 5; i++ { if ok, retry := RateLimitCheckAndRecord(ctx, "probe", window, 5); !ok { t.Fatalf("第 %d 次应放行(窗口内无记录),却 retry=%d", i+1, retry) } } // 放行 5 次后正好到上限,第 6 次该拒,且 retry 很小 ok, retry := RateLimitCheckAndRecord(ctx, "probe", window, 5) if ok { t.Fatal("第 6 次应被拒") } if retry > int(window.Seconds())+2 { t.Errorf("retry_after=%d 超过窗口 %ds", retry, int(window.Seconds())) } } // seedRateLimits 往桶里放 n 条记录,**最早的一条**在 oldestAgo 之前。 // 故意用同一个时间戳:限速器读的是 COUNT 与 MIN(ts) 两个量, // 让它们指向同一刻,测试才只测「retry_after 算得对不对」这���件事。 // rlAbs 取绝对值(包级已有 abs,此处不重名)。 func rlAbs(n int) int { if n < 0 { return -n } return n } func seedRateLimits(t *testing.T, bucket string, n int, oldestAgo time.Duration) { t.Helper() ts := time.Now().Add(-oldestAgo) for i := 0; i < n; i++ { if _, err := db.DB.ExecContext(context.Background(), `INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`, bucket, ts); err != nil { t.Fatalf("seed %d/%d: %v", i+1, n, err) } } }