email_agent/internal/store/triage_test.go

package store

import (
	"context"
	"database/sql"
	"path/filepath"
	"testing"
	"time"
)

func testStore(t *testing.T) (*Store, context.Context) {
	t.Helper()
	ctx := context.Background()
	s, err := Open(ctx, filepath.Join(t.TempDir(), "email_agent.db"))
	if err != nil {
		t.Fatalf("Open() error = %v", err)
	}
	t.Cleanup(func() { _ = s.Close() })
	return s, ctx
}

func seedMessage(t *testing.T, s *Store, ctx context.Context, gmailID, subject, source string) int64 {
	t.Helper()
	id, err := s.UpsertMessage(ctx, &Message{
		GmailID:       gmailID,
		FromAddr:      "sender@example.test",
		ToAddr:        source,
		Subject:       subject,
		SourceAddress: source,
		DateSent:      time.Now().UTC(),
	})
	if err != nil {
		t.Fatalf("seed message: %v", err)
	}
	return id
}

func TestMigrationFromV3AddsFolderColumns(t *testing.T) {
	ctx := context.Background()
	path := filepath.Join(t.TempDir(), "old.db")

	// Simulate a database created by the v3 release: messages has the IMAP
	// columns but not folder/mailbox. Opening it with the new code must apply
	// migration v4 without failing on the folder index (the schema file is
	// executed before migrations).
	db, err := sql.Open("sqlite", "file:"+path)
	if err != nil {
		t.Fatalf("open old db: %v", err)
	}
	stmts := []string{
		`CREATE TABLE messages (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			gmail_id TEXT NOT NULL UNIQUE,
			thread_id TEXT, from_addr TEXT, to_addr TEXT, subject TEXT,
			body_text TEXT, source_address TEXT, date_sent TEXT,
			imap_uid INTEGER, imap_uidvalidity INTEGER,
			received_at TEXT NOT NULL DEFAULT (datetime('now')),
			created_at TEXT NOT NULL DEFAULT (datetime('now')))`,
		`CREATE TABLE schema_migrations (version INTEGER PRIMARY KEY, applied_at TEXT NOT NULL DEFAULT (datetime('now')))`,
		`INSERT INTO schema_migrations (version) VALUES (3)`,
	}
	for _, q := range stmts {
		if _, err := db.ExecContext(ctx, q); err != nil {
			t.Fatalf("seed old db: %v", err)
		}
	}
	_ = db.Close()

	s, err := Open(ctx, path)
	if err != nil {
		t.Fatalf("Open() on v3 database error = %v", err)
	}
	defer s.Close()

	v, err := s.SchemaVersion(ctx)
	if err != nil {
		t.Fatalf("SchemaVersion() error = %v", err)
	}
	if v < 4 {
		t.Errorf("SchemaVersion() = %d, want >= 4", v)
	}
	var n int
	if err := s.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM pragma_table_info('messages') WHERE name = 'folder'`).Scan(&n); err != nil {
		t.Fatalf("count folder column: %v", err)
	}
	if n != 1 {
		t.Errorf("folder column count = %d, want 1", n)
	}
}

func TestMigrationAddsFolderColumns(t *testing.T) {
	s, ctx := testStore(t)

	v, err := s.SchemaVersion(ctx)
	if err != nil {
		t.Fatalf("SchemaVersion() error = %v", err)
	}
	if v < 4 {
		t.Errorf("SchemaVersion() = %d, want >= 4", v)
	}
	for _, col := range []string{"folder", "mailbox"} {
		var found bool
		rows, err := s.DB().QueryContext(ctx, `SELECT name FROM pragma_table_info('messages')`)
		if err != nil {
			t.Fatalf("pragma_table_info: %v", err)
		}
		for rows.Next() {
			var name string
			if err := rows.Scan(&name); err != nil {
				t.Fatalf("scan: %v", err)
			}
			if name == col {
				found = true
			}
		}
		_ = rows.Close()
		if !found {
			t.Errorf("column %q missing after migration", col)
		}
	}
}

func TestUpsertDefaultsToInbox(t *testing.T) {
	s, ctx := testStore(t)
	seedMessage(t, s, ctx, "1", "hello", "a@x.test")

	list, err := s.ListMessagesFiltered(ctx, MessageFilter{Limit: 10})
	if err != nil {
		t.Fatalf("ListMessagesFiltered() error = %v", err)
	}
	if len(list) != 1 {
		t.Fatalf("len = %d, want 1", len(list))
	}
	if list[0].Folder != FolderInbox {
		t.Errorf("folder = %q, want inbox", list[0].Folder)
	}
}

func TestSetLocationAndFolderFilter(t *testing.T) {
	s, ctx := testStore(t)
	idA := seedMessage(t, s, ctx, "1", "keep me", "a@x.test")
	idB := seedMessage(t, s, ctx, "2", "delete me", "b@x.test")

	if err := s.SetMessageLocation(ctx, idB, FolderDelete, "Delete"); err != nil {
		t.Fatalf("SetMessageLocation() error = %v", err)
	}

	del, err := s.ListMessagesFiltered(ctx, MessageFilter{Folder: FolderDelete, Limit: 10})
	if err != nil {
		t.Fatalf("delete filter: %v", err)
	}
	if len(del) != 1 || del[0].ID != idB {
		t.Fatalf("delete filter = %+v, want message %d", del, idB)
	}
	if del[0].Folder != FolderDelete {
		t.Errorf("folder = %q, want delete", del[0].Folder)
	}

	inbox, err := s.ListMessagesFiltered(ctx, MessageFilter{Folder: FolderInbox, Limit: 10})
	if err != nil {
		t.Fatalf("inbox filter: %v", err)
	}
	if len(inbox) != 1 || inbox[0].ID != idA {
		t.Fatalf("inbox filter = %+v, want message %d", inbox, idA)
	}

	all, err := s.ListMessagesFiltered(ctx, MessageFilter{Folder: "all", Limit: 10})
	if err != nil {
		t.Fatalf("all filter: %v", err)
	}
	if len(all) != 2 {
		t.Fatalf("all filter len = %d, want 2", len(all))
	}
}

func TestListMessagesFilteredByCategory(t *testing.T) {
	s, ctx := testStore(t)
	idA := seedMessage(t, s, ctx, "1", "spammy", "a@x.test")
	idB := seedMessage(t, s, ctx, "2", "normal", "b@x.test")
	if _, err := s.SaveClassification(ctx, Classification{MessageID: idA, Category: "Spam", Importance: "high", Action: "delete", Confidence: 0.9}); err != nil {
		t.Fatalf("SaveClassification() error = %v", err)
	}

	spam, err := s.ListMessagesFiltered(ctx, MessageFilter{Category: "spam", Limit: 10})
	if err != nil {
		t.Fatalf("category filter: %v", err)
	}
	if len(spam) != 1 || spam[0].ID != idA {
		t.Fatalf("spam filter = %+v, want message %d", spam, idA)
	}

	unclassified, err := s.ListMessagesFiltered(ctx, MessageFilter{Unclassified: true, Limit: 10})
	if err != nil {
		t.Fatalf("unclassified filter: %v", err)
	}
	if len(unclassified) != 1 || unclassified[0].ID != idB {
		t.Fatalf("unclassified filter = %+v, want message %d", unclassified, idB)
	}
}

func TestLatestClassifications(t *testing.T) {
	s, ctx := testStore(t)
	id := seedMessage(t, s, ctx, "1", "spammy", "a@x.test")
	if _, err := s.SaveClassification(ctx, Classification{MessageID: id, Category: "Info", Action: "keep"}); err != nil {
		t.Fatalf("first classification: %v", err)
	}
	if _, err := s.SaveClassification(ctx, Classification{MessageID: id, Category: "Spam", Action: "delete"}); err != nil {
		t.Fatalf("second classification: %v", err)
	}

	got, err := s.LatestClassifications(ctx, []int64{id})
	if err != nil {
		t.Fatalf("LatestClassifications() error = %v", err)
	}
	if got[id].Category != "Spam" {
		t.Errorf("category = %q, want the newest (Spam)", got[id].Category)
	}
}

func TestCountByFolderAndMessagesByIDs(t *testing.T) {
	s, ctx := testStore(t)
	idA := seedMessage(t, s, ctx, "1", "a", "a@x.test")
	idB := seedMessage(t, s, ctx, "2", "b", "b@x.test")
	if err := s.SetMessageLocation(ctx, idB, FolderSaved, "Saved"); err != nil {
		t.Fatalf("SetMessageLocation() error = %v", err)
	}

	counts, err := s.CountByFolder(ctx)
	if err != nil {
		t.Fatalf("CountByFolder() error = %v", err)
	}
	if counts[FolderInbox] != 1 || counts[FolderSaved] != 1 {
		t.Errorf("counts = %v, want inbox:1 saved:1", counts)
	}

	msgs, err := s.MessagesByIDs(ctx, []int64{idA, idB})
	if err != nil {
		t.Fatalf("MessagesByIDs() error = %v", err)
	}
	if len(msgs) != 2 {
		t.Fatalf("MessagesByIDs() len = %d, want 2", len(msgs))
	}
	for _, m := range msgs {
		if m.ID == idB && (m.Folder != FolderSaved || m.Mailbox != "Saved") {
			t.Errorf("saved message location = (%q, %q)", m.Folder, m.Mailbox)
		}
	}
}