MailUI4Agents/gateway/internal/repo/attachments_test.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))
	}
}