test: add coverage for dedup, collections, identities

Table-driven tests for dedup engine, collections CRUD, identity
discovery, and source filter helpers. Incorporates review findings.

Author: jesserobbins

cmd/msgvault/cmd/collections.go

@@ -191,7 +191,6 @@ func runCollectionsDelete(_ *cobra.Command, args []string) error {
 	return nil
 }
 
-
 func resolveAccountList(st *store.Store, accounts string) ([]int64, error) {
 	if accounts == "" {
 		return nil, fmt.Errorf("--accounts is required")

cmd/msgvault/cmd/deduplicate.go

@@ -62,8 +62,15 @@ func runDeduplicate(cmd *cobra.Command, _ []string) error {
 	preference := dedup.DefaultSourcePreference
 	if dedupPrefer != "" {
 		preference = strings.Split(dedupPrefer, ",")
+		known := make(map[string]bool, len(dedup.DefaultSourcePreference))
+		for _, t := range dedup.DefaultSourcePreference {
+			known[t] = true
+		}
 		for i := range preference {
 			preference[i] = strings.TrimSpace(preference[i])
+			if !known[preference[i]] {
+				fmt.Fprintf(os.Stderr, "Warning: unknown source type in --prefer: %q\n", preference[i])
+			}
 		}
 	}
 
@@ -119,7 +126,7 @@ func runDeduplicate(cmd *cobra.Command, _ []string) error {
 	}
 
 	if len(accountSourceIDs) == 0 {
-		return runDeduplicatePerSource(cmd, st, config)
+		return runDeduplicatePerSource(cmd, st, dbPath, config)
 	}
 
 	return runDeduplicateOnce(cmd, dbPath, config, engine)
@@ -128,6 +135,7 @@ func runDeduplicate(cmd *cobra.Command, _ []string) error {
 func runDeduplicatePerSource(
 	cmd *cobra.Command,
 	st *store.Store,
+	dbPath string,
 	cfgBase dedup.Config,
 ) error {
 	sources, err := st.ListSources("")
@@ -144,6 +152,7 @@ func runDeduplicatePerSource(
 	)
 	fmt.Println()
 
+	backedUp := false
 	anyRan := false
 	for _, src := range sources {
 		cfgScoped := cfgBase
@@ -186,8 +195,21 @@ func runDeduplicatePerSource(
 			}
 		}
 
+		if !backedUp && !dedupNoBackup {
+			backedUp = true
+			backupPath := fmt.Sprintf(
+				"%s.dedup-backup-%s", dbPath,
+				time.Now().Format("20060102-150405"),
+			)
+			fmt.Printf("Backing up database to %s...\n",
+				filepath.Base(backupPath))
+			if err := copyFileForBackup(dbPath, backupPath); err != nil {
+				return fmt.Errorf("backup database: %w", err)
+			}
+		}
+
 		batchID := fmt.Sprintf(
-			"dedup-%s", time.Now().Format("20060102-150405"),
+			"dedup-%s-%d-%s", time.Now().Format("20060102-150405"), src.ID, src.Identifier,
 		)
 		summary, err := engineScoped.Execute(
 			cmd.Context(), report, batchID,
@@ -308,21 +330,21 @@ func readDedupYesNo(cmd *cobra.Command) (bool, error) {
 }
 
 func copyFileForBackup(src, dst string) error {
-	in, err := os.Open(src)
-	if err != nil {
+	// Copy the main database file.
+	if err := copyFile(src, dst); err != nil {
 		return err
 	}
-	defer func() { _ = in.Close() }()
-
-	out, err := os.Create(dst)
-	if err != nil {
-		return err
-	}
-	if _, err := io.Copy(out, in); err != nil {
-		_ = out.Close()
-		return err
+	// Also copy WAL and SHM files if they exist, so the backup is
+	// consistent even when SQLite has uncheckpointed WAL pages.
+	for _, suffix := range []string{"-wal", "-shm"} {
+		extra := src + suffix
+		if _, err := os.Stat(extra); err == nil {
+			if err := copyFile(extra, dst+suffix); err != nil {
+				return fmt.Errorf("copy %s: %w", suffix, err)
+			}
+		}
 	}
-	return out.Close()
+	return nil
 }
 
 func init() {

internal/dedup/dedup.go

@@ -206,7 +206,8 @@ func (e *Engine) Scan(ctx context.Context) (*Report, error) {
 	if count > 0 {
 		e.logger.Info("backfilling rfc822_message_id from stored MIME",
 			"count", count)
-		backfilledCount, err = e.store.BackfillRFC822IDs(
+		var backfillFailed int64
+		backfilledCount, backfillFailed, err = e.store.BackfillRFC822IDs(
 			func(done, total int64) {
 				e.logger.Info("backfill progress",
 					"done", done, "total", total)
@@ -219,6 +220,10 @@ func (e *Engine) Scan(ctx context.Context) (*Report, error) {
 			e.logger.Info("backfilled rfc822_message_id",
 				"count", backfilledCount)
 		}
+		if backfillFailed > 0 {
+			e.logger.Warn("backfill: some messages could not be parsed",
+				"failed", backfillFailed)
+		}
 	}
 
 	totalMessages, err := e.store.CountActiveMessages(
@@ -517,7 +522,14 @@ func normalizeRawMIME(raw []byte) []byte {
 	}
 
 	headerSection := raw[:headerEnd]
-	body := raw[headerEnd:]
+	// Find the start of the actual body after the blank line.
+	bodyStart := headerEnd
+	if bytes.HasPrefix(raw[headerEnd:], []byte("\r\n\r\n")) {
+		bodyStart = headerEnd + 4
+	} else {
+		bodyStart = headerEnd + 2 // "\n\n"
+	}
+	body := raw[bodyStart:]
 
 	// Copy headerSection before appending to avoid mutating the
 	// underlying raw buffer (headerSection is a sub-slice of raw).
@@ -544,6 +556,7 @@ func normalizeRawMIME(raw []byte) []byte {
 			fmt.Fprintf(&buf, "%s: %s\n", key, val)
 		}
 	}
+	buf.WriteString("\n") // canonical header/body separator
 	buf.Write(body)
 	return buf.Bytes()
 }

internal/dedup/dedup_test.go

@@ -364,3 +364,84 @@ func TestEngine_FormatMethodology_MentionsSentPolicy(t *testing.T) {
 		t.Errorf("methodology missing cross-account guarantee")
 	}
 }
+
+func TestEngine_SurvivorTiebreakers(t *testing.T) {
+	t.Run("raw MIME wins over no raw MIME", func(t *testing.T) {
+		f := storetest.New(t)
+		st := f.Store
+
+		idNoRaw := addMessage(t, st, f.Source, "no-raw", "rfc-raw-tie", false)
+		idHasRaw := addMessage(t, st, f.Source, "has-raw", "rfc-raw-tie", false)
+		testutil.MustNoErr(t,
+			st.UpsertMessageRaw(idHasRaw, []byte("Subject: test\r\n\r\nBody")),
+			"UpsertMessageRaw",
+		)
+
+		eng := dedup.NewEngine(st, dedup.Config{
+			AccountSourceIDs: []int64{f.Source.ID},
+			Account:          "test",
+		}, nil)
+		report, err := eng.Scan(context.Background())
+		testutil.MustNoErr(t, err, "Scan")
+		if report.DuplicateGroups != 1 {
+			t.Fatalf("groups = %d, want 1", report.DuplicateGroups)
+		}
+		survivor := report.Groups[0].Messages[report.Groups[0].Survivor]
+		if survivor.ID != idHasRaw {
+			t.Errorf("survivor = %d, want %d (has raw)", survivor.ID, idHasRaw)
+		}
+		_ = idNoRaw
+	})
+
+	t.Run("more labels wins when raw MIME is equal", func(t *testing.T) {
+		f := storetest.New(t)
+		st := f.Store
+
+		idFew := addMessage(t, st, f.Source, "few", "rfc-label-tie", false)
+		idMany := addMessage(t, st, f.Source, "many", "rfc-label-tie", false)
+
+		lid1, _ := st.EnsureLabel(f.Source.ID, "L1", "Label1", "user")
+		lid2, _ := st.EnsureLabel(f.Source.ID, "L2", "Label2", "user")
+		lid3, _ := st.EnsureLabel(f.Source.ID, "L3", "Label3", "user")
+		_ = st.LinkMessageLabel(idFew, lid1)
+		_ = st.LinkMessageLabel(idMany, lid1)
+		_ = st.LinkMessageLabel(idMany, lid2)
+		_ = st.LinkMessageLabel(idMany, lid3)
+
+		eng := dedup.NewEngine(st, dedup.Config{
+			AccountSourceIDs: []int64{f.Source.ID},
+			Account:          "test",
+		}, nil)
+		report, err := eng.Scan(context.Background())
+		testutil.MustNoErr(t, err, "Scan")
+		if report.DuplicateGroups != 1 {
+			t.Fatalf("groups = %d, want 1", report.DuplicateGroups)
+		}
+		survivor := report.Groups[0].Messages[report.Groups[0].Survivor]
+		if survivor.ID != idMany {
+			t.Errorf("survivor = %d, want %d (more labels)", survivor.ID, idMany)
+		}
+	})
+
+	t.Run("lower ID wins as final tiebreaker", func(t *testing.T) {
+		f := storetest.New(t)
+		st := f.Store
+
+		idFirst := addMessage(t, st, f.Source, "first", "rfc-id-tie", false)
+		_ = addMessage(t, st, f.Source, "second", "rfc-id-tie", false)
+
+		eng := dedup.NewEngine(st, dedup.Config{
+			AccountSourceIDs: []int64{f.Source.ID},
+			Account:          "test",
+		}, nil)
+		report, err := eng.Scan(context.Background())
+		testutil.MustNoErr(t, err, "Scan")
+		if report.DuplicateGroups != 1 {
+			t.Fatalf("groups = %d, want 1", report.DuplicateGroups)
+		}
+		survivor := report.Groups[0].Messages[report.Groups[0].Survivor]
+		if survivor.ID != idFirst {
+			t.Errorf("survivor = %d, want %d (lower ID)", survivor.ID, idFirst)
+		}
+	})
+}

internal/dedup/normalize_test.go

@@ -0,0 +1,92 @@
+package dedup
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestNormalizeRawMIME(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    []byte
+		wantSame bool   // true if output should equal input
+		contains string // substring the output must contain
+		excludes string // substring the output must NOT contain
+	}{
+		{
+			name:     "strips Received header (CRLF)",
+			input:    []byte("Received: from mx1.google.com\r\nFrom: alice@example.com\r\nSubject: Hi\r\n\r\nBody"),
+			contains: "From: alice@example.com",
+			excludes: "Received",
+		},
+		{
+			name:     "strips multiple transport headers",
+			input:    []byte("Delivered-To: bob@example.com\r\nX-Gmail-Labels: INBOX\r\nAuthentication-Results: spf=pass\r\nFrom: alice@example.com\r\nSubject: Test\r\n\r\nBody"),
+			contains: "From: alice@example.com",
+			excludes: "Delivered-To",
+		},
+		{
+			name:     "preserves non-transport headers",
+			input:    []byte("From: alice@example.com\r\nTo: bob@example.com\r\nSubject: Meeting\r\nDate: Mon, 1 Jan 2024 12:00:00 +0000\r\n\r\nBody text"),
+			contains: "Subject: Meeting",
+		},
+		{
+			name:     "handles LF-only line endings",
+			input:    []byte("Received: from mx1\nFrom: alice@example.com\nSubject: Test\n\nBody with LF"),
+			contains: "From: alice@example.com",
+			excludes: "Received",
+		},
+		{
+			name:     "no header/body separator returns raw unchanged",
+			input:    []byte("This is just a blob of text with no headers"),
+			wantSame: true,
+		},
+		{
+			name:     "empty body preserved",
+			input:    []byte("From: alice@example.com\r\nSubject: Empty\r\n\r\n"),
+			contains: "Subject: Empty",
+		},
+		{
+			name:     "preserves body content exactly",
+			input:    []byte("Received: from mx1\r\nFrom: a@b.com\r\n\r\nExact body content here."),
+			contains: "Exact body content here.",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			inputCopy := make([]byte, len(tt.input))
+			copy(inputCopy, tt.input)
+
+			result := normalizeRawMIME(tt.input)
+
+			if !bytes.Equal(tt.input, inputCopy) {
+				t.Error("normalizeRawMIME mutated its input buffer")
+			}
+
+			if tt.wantSame {
+				if !bytes.Equal(result, tt.input) {
+					t.Errorf("expected unchanged output, got:\n%s", result)
+				}
+				return
+			}
+			if tt.contains != "" && !bytes.Contains(result, []byte(tt.contains)) {
+				t.Errorf("output missing %q:\n%s", tt.contains, result)
+			}
+			if tt.excludes != "" && bytes.Contains(result, []byte(tt.excludes)) {
+				t.Errorf("output should not contain %q:\n%s", tt.excludes, result)
+			}
+		})
+	}
+}
+
+func TestNormalizeRawMIME_DeterministicOutput(t *testing.T) {
+	raw1 := []byte("Received: from mx1.google.com\r\nFrom: sender@example.com\r\nSubject: Meeting\r\nDate: Mon, 1 Jan 2024 12:00:00 +0000\r\n\r\nLet's meet at 3pm.")
+	raw2 := []byte("Received: from mx2.google.com\r\nDelivered-To: other@example.com\r\nFrom: sender@example.com\r\nSubject: Meeting\r\nDate: Mon, 1 Jan 2024 12:00:00 +0000\r\n\r\nLet's meet at 3pm.")
+
+	hash1 := sha256Hex(normalizeRawMIME(raw1))
+	hash2 := sha256Hex(normalizeRawMIME(raw2))
+	if hash1 != hash2 {
+		t.Errorf("same message with different transport headers produced different hashes")
+	}
+}

internal/query/duckdb.go

@@ -682,10 +682,7 @@ func (e *DuckDBEngine) buildWhereClause(opts AggregateOptions, keyColumns ...str
 	// message_type IS NULL and '' handle old data without the column.
 	conditions = append(conditions, "(msg.message_type = 'email' OR msg.message_type IS NULL OR msg.message_type = '')")
 
-	if opts.SourceID != nil {
-		conditions = append(conditions, "msg.source_id = ?")
-		args = append(args, *opts.SourceID)
-	}
+	conditions, args = appendSourceFilter(conditions, args, "msg.", opts.SourceID, opts.SourceIDs)
 
 	if opts.After != nil {
 		conditions = append(conditions, "msg.sent_at >= CAST(? AS TIMESTAMP)")
@@ -888,10 +885,7 @@ func (e *DuckDBEngine) buildFilterConditions(filter MessageFilter) (string, []in
 	// message_type IS NULL and '' handle old data without the column.
 	conditions = append(conditions, "(msg.message_type = 'email' OR msg.message_type IS NULL OR msg.message_type = '')")
 
-	if filter.SourceID != nil {
-		conditions = append(conditions, "msg.source_id = ?")
-		args = append(args, *filter.SourceID)
-	}
+	conditions, args = appendSourceFilter(conditions, args, "msg.", filter.SourceID, filter.SourceIDs)
 
 	if filter.ConversationID != nil {
 		conditions = append(conditions, "msg.conversation_id = ?")
@@ -1151,10 +1145,7 @@ func (e *DuckDBEngine) GetTotalStats(ctx context.Context, opts StatsOptions) (*T
 	// Restrict to email messages only; NULL and '' handle pre-message_type data.
 	conditions = append(conditions, emailOnlyFilterMsg)
 
-	if opts.SourceID != nil {
-		conditions = append(conditions, "msg.source_id = ?")
-		args = append(args, *opts.SourceID)
-	}
+	conditions, args = appendSourceFilter(conditions, args, "msg.", opts.SourceID, opts.SourceIDs)
 
 	if opts.WithAttachmentsOnly {
 		conditions = append(conditions, "msg.has_attachments = 1")
@@ -2343,10 +2334,7 @@ func (e *DuckDBEngine) buildSearchConditions(q *search.Query, filter MessageFilt
 	conditions = append(conditions, emailOnlyFilterMsg)
 
 	// Apply basic filter conditions (ignoring join flags for search - we handle those differently)
-	if filter.SourceID != nil {
-		conditions = append(conditions, "msg.source_id = ?")
-		args = append(args, *filter.SourceID)
-	}
+	conditions, args = appendSourceFilter(conditions, args, "msg.", filter.SourceID, filter.SourceIDs)
 	if filter.After != nil {
 		conditions = append(conditions, "msg.sent_at >= CAST(? AS TIMESTAMP)")
 		args = append(args, filter.After.Format("2006-01-02 15:04:05"))