126 lines
3.5 KiB
Go
126 lines
3.5 KiB
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())
|
||
}
|