MailUI4Agents/gateway/internal/repo/quota_test.go

package repo

import (
	"context"
	"os"
	"path/filepath"
	"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, default_rounds) VALUES ($1, 'x', 'test', $2)`,
		name, max)
	if err != nil {
		t.Fatalf("seed agent: %v", err)
	}
}

// default_rounds 是「派给这个 Agent 的新任务默认多少个来回」，
// 不是会拦请求的终身额度 —— 真正的额度在 sessions.max_rounds 上。
func TestDefaultRoundsRoundTrip(t *testing.T) {
	setupTestDB(t)
	seedAgent(t, "bot", 0)
	ctx := context.Background()

	st, err := SetDefaultRounds(ctx, "bot", 15)
	if err != nil {
		t.Fatal(err)
	}
	if st.DefaultRounds != 15 {
		t.Fatalf("default_rounds = %d，期望 15", st.DefaultRounds)
	}
	if got := DefaultRoundsFor(ctx, "bot"); got != 15 {
		t.Fatalf("DefaultRoundsFor = %d，期望 15", got)
	}

	// 负数归一为 0（不限），而不是造出一个永远发不出信的默认值
	if st, err = SetDefaultRounds(ctx, "bot", -3); err != nil {
		t.Fatal(err)
	}
	if st.DefaultRounds != 0 {
		t.Fatalf("负数应归一为 0，实际 %d", st.DefaultRounds)
	}
}

// 未注册的 Agent 取默认预算时给兜底值而不是报错：
// 派活的人不该因为「对方还没上线」就拿不到一个合理默认值 ——
// 邮件本来就支持发给尚未上线的收件人。
func TestDefaultRoundsForUnknownAgentFallsBack(t *testing.T) {
	setupTestDB(t)
	if got := DefaultRoundsFor(context.Background(), "ghost"); got != fallbackDefaultRounds {
		t.Fatalf("未注册 Agent 应回落到 %d，实际 %d", fallbackDefaultRounds, got)
	}
}

// BumpSentCount 是纯统计：只累加，绝不拦请求，也绝不返回错误
func TestBumpSentCountOnlyCounts(t *testing.T) {
	setupTestDB(t)
	seedAgent(t, "bot", 0)
	ctx := context.Background()

	for i := 0; i < 5; i++ {
		BumpSentCount(ctx, "bot")
	}
	st, err := GetAgentStats(ctx, "bot")
	if err != nil {
		t.Fatal(err)
	}
	if st.SentTotal != 5 {
		t.Fatalf("累计发信 = %d，期望 5", st.SentTotal)
	}

	// 不存在的 Agent 也不该 panic 或报错 —— 它只是没有行可更新
	BumpSentCount(ctx, "ghost")
}

func TestGetAgentStatsUnknownAgent(t *testing.T) {
	setupTestDB(t)
	if _, err := GetAgentStats(context.Background(), "ghost"); err == nil {
		t.Fatal("不存在的 Agent 应报错")
	}
}

func TestListAgentStats(t *testing.T) {
	setupTestDB(t)
	seedAgent(t, "alpha", 0)
	seedAgent(t, "beta", 0)
	ctx := context.Background()
	SetDefaultRounds(ctx, "alpha", 5)

	list, err := ListAgentStats(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if len(list) != 2 {
		t.Fatalf("应有 2 个 Agent，实际 %d", len(list))
	}
	// 按名字排序，alpha 在前
	if list[0].AgentName != "alpha" || list[0].DefaultRounds != 5 {
		t.Fatalf("alpha 的记录不对：%+v", list[0])
	}
}

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