Maven/internal/store/backfill_test.go

package store

import (
	"context"
	"errors"
	"testing"
	"time"
)

// markerVec is a recognisable vector: nothing in these tests writes it except
// the backfill, so finding it proves the row really was rewritten.
var markerVec = []float32{9, 9, 9}

func newEmbedder(calls *int) EmbedFunc {
	return func(_ context.Context, _ string) ([]float32, error) {
		*calls++
		return markerVec, nil
	}
}

// seedOldVectors puts one note (notes table + unified index) and one fact
// (unified index only) in the DB, both carrying obviously-old vectors.
func seedOldVectors(t *testing.T, s *Store) {
	t.Helper()
	ctx := context.Background()
	old := []float32{0.1, 0.2, 0.3}
	id, err := s.WriteNote(ctx, time.Now(), "молоко в холодильнике", old, "voice")
	if err != nil {
		t.Fatalf("WriteNote: %v", err)
	}
	mem := s.VectorMemory()
	if err := mem.Insert(ctx, "note:1", old, map[string]string{
		"type": "note", "text": "молоко в холодильнике",
	}); err != nil {
		t.Fatalf("Insert note vector: %v", err)
	}
	if err := mem.Insert(ctx, "fact:water:1", old, map[string]string{
		"type": "fact", "text": "я пил воду",
	}); err != nil {
		t.Fatalf("Insert fact vector: %v", err)
	}
	_ = id
}

func noteVec(t *testing.T, s *Store) []float32 {
	t.Helper()
	var blob []byte
	if err := s.db.QueryRow(`SELECT embedding FROM notes LIMIT 1`).Scan(&blob); err != nil {
		t.Fatalf("read note embedding: %v", err)
	}
	return blobToFloats(blob)
}

func memVec(t *testing.T, s *Store, id string) []float32 {
	t.Helper()
	var blob []byte
	if err := s.db.QueryRow(`SELECT vec FROM memory_vectors WHERE id = ?`, id).Scan(&blob); err != nil {
		t.Fatalf("read memory vector %s: %v", id, err)
	}
	return decodeVec(blob)
}

func sameVec(a, b []float32) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// The deployed case: old vectors everywhere, no marker. Every vector in both
// places must be rewritten and the marker recorded.
func TestReembedAllRewritesEveryVector(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	seedOldVectors(t, s)

	calls := 0
	res, err := s.ReembedAll(ctx, "multilingual-e5-small@384", newEmbedder(&calls))
	if err != nil {
		t.Fatalf("ReembedAll: %v", err)
	}
	if res.Skipped {
		t.Fatal("first run should not skip")
	}
	if res.Notes != 1 || res.MemNotes != 1 || res.Facts != 1 {
		t.Fatalf("counts: notes=%d memNotes=%d facts=%d", res.Notes, res.MemNotes, res.Facts)
	}
	if calls != 3 {
		t.Fatalf("embedder called %d times, want 3", calls)
	}
	if !sameVec(noteVec(t, s), markerVec) {
		t.Fatalf("notes table not rewritten: %v", noteVec(t, s))
	}
	if !sameVec(memVec(t, s, "note:1"), markerVec) {
		t.Fatal("unified index note row not rewritten")
	}
	if !sameVec(memVec(t, s, "fact:water:1"), markerVec) {
		t.Fatal("unified index fact row not rewritten")
	}
	got, err := s.Meta(ctx, metaKeyEmbedderID)
	if err != nil {
		t.Fatalf("Meta: %v", err)
	}
	if got != "multilingual-e5-small@384" {
		t.Fatalf("marker = %q", got)
	}
}

// Re-running must do nothing at all — not a second pass over the same rows.
func TestReembedAllSecondRunIsNoop(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	seedOldVectors(t, s)

	calls := 0
	if _, err := s.ReembedAll(ctx, "e5@384", newEmbedder(&calls)); err != nil {
		t.Fatalf("first run: %v", err)
	}
	first := calls

	res, err := s.ReembedAll(ctx, "e5@384", newEmbedder(&calls))
	if err != nil {
		t.Fatalf("second run: %v", err)
	}
	if !res.Skipped {
		t.Fatal("second run should report Skipped")
	}
	if calls != first {
		t.Fatalf("second run embedded %d more rows, want 0", calls-first)
	}
}

// A failure partway must leave the DB exactly as it was: no marker, and the old
// vectors still in place (one transaction, rolled back).
func TestReembedAllPartialFailureLeavesMarkerUnset(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	seedOldVectors(t, s)
	before := noteVec(t, s)

	calls := 0
	boom := func(_ context.Context, _ string) ([]float32, error) {
		calls++
		if calls == 2 {
			return nil, errors.New("onnx blew up")
		}
		return markerVec, nil
	}
	if _, err := s.ReembedAll(ctx, "e5@384", boom); err == nil {
		t.Fatal("expected an error")
	}
	got, err := s.Meta(ctx, metaKeyEmbedderID)
	if err != nil {
		t.Fatalf("Meta: %v", err)
	}
	if got != "" {
		t.Fatalf("marker was set to %q after a failed run", got)
	}
	if !sameVec(noteVec(t, s), before) {
		t.Fatal("a failed run left a partially rewritten notes table")
	}
	// And the mismatch warning must still fire, so the user knows to re-run.
	if _, mismatch, err := s.CheckEmbedder(ctx, "e5@384"); err != nil || !mismatch {
		t.Fatalf("CheckEmbedder after failed backfill: mismatch=%v err=%v", mismatch, err)
	}
}