chore: directory migration - gateway→server, web→client/electron

server/internal/repo/quota_test.go

@@ -0,0 +1,125 @@
+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())
+}