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()) }