MailUI4Agents/gateway/internal/repo/quota_test.go

package repo

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"sync"
	"testing"

	"github.com/agentmail/gateway/internal/db"
)

// setupTestDB 起一个临时 SQLite 库并建表，供配额测试使用。
// 直接用真实的 SQLite 而非 mock：配额的正确性核心在于「判断与自增在同一条 UPDATE 里」，
// 这正是只有真实数据库才能验证的部分。
func setupTestDB(t *testing.T) {
	t.Helper()

	dir := t.TempDir()
	if err := db.Connect(context.Background(), filepath.Join(dir, "test.db")); err != nil {
		t.Fatalf("connect: %v", err)
	}
	if err := db.Migrate(context.Background()); err != nil {
		t.Fatalf("migrate: %v", err)
	}
	t.Cleanup(db.Close)
}

func seedAgent(t *testing.T, name string, max int) {
	t.Helper()
	_, err := db.DB.ExecContext(context.Background(),
		`INSERT INTO agents (agent_name, secret, platform, max_rounds) VALUES ($1, 'x', 'test', $2)`,
		name, max)
	if err != nil {
		t.Fatalf("seed agent: %v", err)
	}
}

func TestConsumeQuotaCountsDown(t *testing.T) {
	setupTestDB(t)
	seedAgent(t, "bot", 3)
	ctx := context.Background()

	for i := 1; i <= 3; i++ {
		q, err := ConsumeQuota(ctx, "bot")
		if err != nil {
			t.Fatalf("第 %d 次占用失败: %v", i, err)
		}
		if q.Used != i || q.Remaining != 3-i {
			t.Errorf("第 %d 次: used=%d remaining=%d, want used=%d remaining=%d",
				i, q.Used, q.Remaining, i, 3-i)
		}
	}

	q, err := ConsumeQuota(ctx, "bot")
	if !errors.Is(err, ErrQuotaExhausted) {
		t.Fatalf("第 4 次应耗尽，得到 err=%v", err)
	}
	// 耗尽时仍要给出快照，调用方才能在错误文案里写清 used/max
	if q.Used != 3 || q.Max != 3 {
		t.Errorf("耗尽时快照不对: %+v", q)
	}
}

func TestConsumeQuotaUnlimited(t *testing.T) {
	setupTestDB(t)
	seedAgent(t, "free", 0) // 0 = 不限
	ctx := context.Background()

	for i := 0; i < 5; i++ {
		q, err := ConsumeQuota(ctx, "free")
		if err != nil {
			t.Fatalf("不限额时不应失败: %v", err)
		}
		if !q.Unlimited || q.Remaining != -1 {
			t.Errorf("不限额快照不对: %+v", q)
		}
	}
}

// 配额的核心保证：并发发信不能把额度刷穿。
// 判断与自增若分成两步（先 SELECT 再 UPDATE），并发下两个请求会双双通过检查。
func TestConsumeQuotaConcurrentDoesNotOverdraw(t *testing.T) {
	setupTestDB(t)
	const limit = 10
	seedAgent(t, "racer", limit)
	ctx := context.Background()

	const attempts = 40
	var (
		wg      sync.WaitGroup
		mu      sync.Mutex
		granted int
	)
	for i := 0; i < attempts; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if _, err := ConsumeQuota(ctx, "racer"); err == nil {
				mu.Lock()
				granted++
				mu.Unlock()
			}
		}()
	}
	wg.Wait()

	if granted != limit {
		t.Errorf("并发 %d 次请求在上限 %d 下放行了 %d 次", attempts, limit, granted)
	}

	q, err := GetQuota(ctx, "racer")
	if err != nil {
		t.Fatal(err)
	}
	if q.Used != limit {
		t.Errorf("used_rounds = %d，应恰好等于上限 %d（未刷穿也未少记）", q.Used, limit)
	}
}

func TestSetAndResetQuota(t *testing.T) {
	setupTestDB(t)
	seedAgent(t, "bot", 2)
	ctx := context.Background()

	if _, err := ConsumeQuota(ctx, "bot"); err != nil {
		t.Fatal(err)
	}

	q, err := SetQuota(ctx, "bot", 5)
	if err != nil {
		t.Fatal(err)
	}
	// 调上限不应清掉已用次数
	if q.Max != 5 || q.Used != 1 || q.Remaining != 4 {
		t.Errorf("SetQuota 后 %+v，want max=5 used=1 remaining=4", q)
	}

	q, err = ResetQuota(ctx, "bot")
	if err != nil {
		t.Fatal(err)
	}
	if q.Used != 0 || q.Remaining != 5 {
		t.Errorf("ResetQuota 后 %+v，want used=0 remaining=5", q)
	}

	// 负数上限归一为 0（不限），而不是造出一个永远发不出信的 Agent
	if q, err = SetQuota(ctx, "bot", -3); err != nil {
		t.Fatal(err)
	} else if !q.Unlimited {
		t.Errorf("负数上限应视为不限，得到 %+v", q)
	}
}

func TestQuotaUnknownAgent(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()

	if _, err := GetQuota(ctx, "ghost"); err == nil {
		t.Error("不存在的 Agent 应报错")
	}
	if _, err := ConsumeQuota(ctx, "ghost"); err == nil {
		t.Error("不存在的 Agent 占用配额应报错")
	}
	if errors.Is(func() error { _, e := ConsumeQuota(ctx, "ghost"); return e }(), ErrQuotaExhausted) {
		t.Error("不存在的 Agent 不该被报成「配额耗尽」")
	}
}

func TestMain(m *testing.M) {
	os.Exit(m.Run())
}