MailUI4Agents/server/internal/repo/ratelimit_test.go

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 算得对不对」这<E3808D><E8BF99><EFBFBD>件事。
// 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)
		}
	}
}