143 lines
4.5 KiB
Go
143 lines
4.5 KiB
Go
package repo
|
||
|
||
import (
|
||
"context"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/agentmail/gateway/internal/db"
|
||
"github.com/google/uuid"
|
||
)
|
||
|
||
// seedClock 给测试数据发严格递增的时间戳。
|
||
//
|
||
// 不靠挂钟：测试在一个循环里连插几封，很可能落在同一毫秒里，
|
||
// 于是「会话里最早/最后那封」的排序由 mail_id（随机 UUID）决定 —— 结果随机。
|
||
// 生产里两封邮件至少隔着一次模型推理，同毫秒撞车不现实；
|
||
// 但测试必须确定，所以显式发号。
|
||
var seedClock = time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
||
|
||
func nextSeedTime() string {
|
||
seedClock = seedClock.Add(time.Second)
|
||
return seedClock.Format("2006-01-02 15:04:05.000")
|
||
}
|
||
|
||
// seedMailIn 在指定会话里插一封邮件，返回其 id。
|
||
// 时间戳严格递增，因此调用顺序就是邮件的先后顺序。
|
||
func seedMailIn(t *testing.T, sessionID uuid.UUID, from, to, subject string) uuid.UUID {
|
||
t.Helper()
|
||
var id uuid.UUID
|
||
err := db.DB.QueryRowContext(context.Background(), `
|
||
INSERT INTO mails (session_id, from_name, from_workspace, to_name, to_workspace,
|
||
subject, body, cc_list, created_at)
|
||
VALUES ($1, $2, '', $3, '', $4, 'body', '[]', $5)
|
||
RETURNING mail_id
|
||
`, sessionID, from, to, subject, nextSeedTime()).Scan(&id)
|
||
if err != nil {
|
||
t.Fatalf("seed mail: %v", err)
|
||
}
|
||
return id
|
||
}
|
||
|
||
func seedSessionRow(t *testing.T, alias string) uuid.UUID {
|
||
t.Helper()
|
||
var id uuid.UUID
|
||
err := db.DB.QueryRowContext(context.Background(), `
|
||
INSERT INTO sessions (from_agent, subject, session_alias)
|
||
VALUES ('opencode', 'attach test', $1)
|
||
RETURNING session_id
|
||
`, alias).Scan(&id)
|
||
if err != nil {
|
||
t.Fatalf("seed session: %v", err)
|
||
}
|
||
return id
|
||
}
|
||
|
||
func attach(t *testing.T, mailID uuid.UUID, name, sum string) {
|
||
t.Helper()
|
||
a, err := CreateAttachment(context.Background(), "admin", name, "text/plain", 3, sum)
|
||
if err != nil {
|
||
t.Fatalf("create attachment: %v", err)
|
||
}
|
||
if err := AttachToMail(context.Background(), mailID, []uuid.UUID{a.ID}, "admin"); err != nil {
|
||
t.Fatalf("attach: %v", err)
|
||
}
|
||
}
|
||
|
||
// ListAttachmentsForMails 存在的理由是消掉 N+1：
|
||
// 原先每封邮件单独查一次，一个 200 封的会话打开要打 200 次库。
|
||
func TestListAttachmentsForMails(t *testing.T) {
|
||
setupTestDB(t)
|
||
sid := seedSessionRow(t, "batch-attach")
|
||
|
||
m1 := seedMailIn(t, sid, "admin", "opencode", "两个附件")
|
||
m2 := seedMailIn(t, sid, "opencode", "admin", "一个附件")
|
||
m3 := seedMailIn(t, sid, "admin", "opencode", "没有附件")
|
||
|
||
attach(t, m1, "a.txt", "sum-a")
|
||
attach(t, m1, "b.txt", "sum-b")
|
||
attach(t, m2, "c.txt", "sum-c")
|
||
|
||
got, err := ListAttachmentsForMails(context.Background(),
|
||
[]uuid.UUID{m1, m2, m3})
|
||
if err != nil {
|
||
t.Fatalf("批量查询失败: %v", err)
|
||
}
|
||
|
||
if n := len(got[m1]); n != 2 {
|
||
t.Errorf("m1 应有 2 个附件，实际 %d", n)
|
||
}
|
||
if n := len(got[m2]); n != 1 {
|
||
t.Errorf("m2 应有 1 个附件，实际 %d", n)
|
||
}
|
||
// 无附件的邮件不该出现在 map 里：调用方据此保持 Attachments 为 nil，
|
||
// 这样带 omitempty 的字段不会给每封邮件的 JSON 白加一个 "attachments":[]
|
||
if _, ok := got[m3]; ok {
|
||
t.Errorf("m3 无附件却出现在结果里：%#v", got[m3])
|
||
}
|
||
|
||
// 同一封内按 created_at 排序，顺序不能乱
|
||
if len(got[m1]) == 2 && got[m1][0].Filename != "a.txt" {
|
||
t.Errorf("同一封内应按上传顺序，首个是 %s", got[m1][0].Filename)
|
||
}
|
||
|
||
// 每条都要带回 mail_id，否则调用方分不清是谁的
|
||
for _, a := range got[m1] {
|
||
if a.MailID == nil || *a.MailID != m1 {
|
||
t.Errorf("附件 %s 的 mail_id 不对：%v", a.Filename, a.MailID)
|
||
}
|
||
}
|
||
}
|
||
|
||
// 空输入必须返回空 map 而非 nil：调用方直接索引不该 panic。
|
||
func TestListAttachmentsForMailsEmpty(t *testing.T) {
|
||
setupTestDB(t)
|
||
|
||
for _, ids := range [][]uuid.UUID{nil, {}} {
|
||
got, err := ListAttachmentsForMails(context.Background(), ids)
|
||
if err != nil {
|
||
t.Fatalf("空输入不该报错: %v", err)
|
||
}
|
||
if got == nil {
|
||
t.Fatal("空输入应返回空 map 而非 nil")
|
||
}
|
||
if len(got) != 0 {
|
||
t.Errorf("空输入应返回空结果，实际 %d 项", len(got))
|
||
}
|
||
}
|
||
}
|
||
|
||
// 传入不存在的 mail_id 不该报错，只是查不到 —— 调用方可能拿着已删邮件的 id。
|
||
func TestListAttachmentsForMailsUnknownID(t *testing.T) {
|
||
setupTestDB(t)
|
||
|
||
got, err := ListAttachmentsForMails(context.Background(),
|
||
[]uuid.UUID{uuid.New(), uuid.New()})
|
||
if err != nil {
|
||
t.Fatalf("未知 id 不该报错: %v", err)
|
||
}
|
||
if len(got) != 0 {
|
||
t.Errorf("未知 id 应查不到，实际 %d 项", len(got))
|
||
}
|
||
}
|