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