golangLAKEHOUSE/internal/vectord/index_test.go

package vectord

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"sync"
	"testing"
)

func TestNewIndex_DefaultsAndValidation(t *testing.T) {
	idx, err := NewIndex(IndexParams{Name: "x", Dimension: 4})
	if err != nil {
		t.Fatal(err)
	}
	p := idx.Params()
	if p.M != DefaultM || p.EfSearch != DefaultEfSearch || p.Distance != DistanceCosine {
		t.Errorf("defaults not applied: %+v", p)
	}

	if _, err := NewIndex(IndexParams{Dimension: 4}); err == nil {
		t.Error("expected error on empty name")
	}
	if _, err := NewIndex(IndexParams{Name: "y", Dimension: 0}); err == nil {
		t.Error("expected error on zero dimension")
	}
	if _, err := NewIndex(IndexParams{Name: "z", Dimension: 4, Distance: "bogus"}); !errors.Is(err, ErrUnknownDistance) {
		t.Errorf("expected ErrUnknownDistance, got %v", err)
	}
}

func TestIndex_AddAndSearch_Recall(t *testing.T) {
	idx, err := NewIndex(IndexParams{Name: "x", Dimension: 3, Distance: DistanceCosine})
	if err != nil {
		t.Fatal(err)
	}
	target := []float32{1, 0, 0}
	if err := idx.Add("alice", target, json.RawMessage(`{"role":"warehouse"}`)); err != nil {
		t.Fatal(err)
	}
	if err := idx.Add("bob", []float32{0, 1, 0}, nil); err != nil {
		t.Fatal(err)
	}
	if err := idx.Add("carol", []float32{0, 0, 1}, nil); err != nil {
		t.Fatal(err)
	}
	if idx.Len() != 3 {
		t.Errorf("Len: got %d, want 3", idx.Len())
	}
	hits, err := idx.Search(target, 2)
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) < 1 {
		t.Fatal("no hits")
	}
	if hits[0].ID != "alice" {
		t.Errorf("nearest: got %q, want alice", hits[0].ID)
	}
	if hits[0].Distance > 0.001 {
		t.Errorf("nearest distance: got %v, want ~0", hits[0].Distance)
	}
	// Cosine distance of identical unit vectors is 0; metadata round-trips.
	if string(hits[0].Metadata) != `{"role":"warehouse"}` {
		t.Errorf("metadata: got %q", hits[0].Metadata)
	}
}

func TestIndex_DimensionMismatch(t *testing.T) {
	idx, _ := NewIndex(IndexParams{Name: "x", Dimension: 4})
	err := idx.Add("a", []float32{1, 2, 3}, nil)
	if !errors.Is(err, ErrDimensionMismatch) {
		t.Errorf("Add: expected ErrDimensionMismatch, got %v", err)
	}
	_, err = idx.Search([]float32{1, 2, 3}, 1)
	if !errors.Is(err, ErrDimensionMismatch) {
		t.Errorf("Search: expected ErrDimensionMismatch, got %v", err)
	}
}

func TestIndex_DeleteAndLookup(t *testing.T) {
	idx, _ := NewIndex(IndexParams{Name: "x", Dimension: 2})
	_ = idx.Add("a", []float32{1, 0}, nil)
	if !idx.Delete("a") {
		t.Error("Delete returned false on existing key")
	}
	if _, _, ok := idx.Lookup("a"); ok {
		t.Error("Lookup found deleted key")
	}
	if idx.Delete("a") {
		t.Error("Delete should return false on missing key")
	}
}

// TestIndex_ConcurrentSearchAdd exercises the RWMutex — many
// concurrent searches alongside a writer adding distinct keys
// shouldn't deadlock, panic, or interleave incorrectly. Each
// writer goroutine gets its own key namespace so we don't
// stress-test the library's re-add path (which has known issues
// under high churn — the wrapper exposes idempotent semantics
// via single-threaded Delete+Add but isn't a fix-everything for
// upstream).
func TestIndex_ConcurrentSearchAdd(t *testing.T) {
	idx, _ := NewIndex(IndexParams{Name: "x", Dimension: 4})
	for i := 0; i < 50; i++ {
		_ = idx.Add(fmt.Sprintf("seed-%d", i), []float32{float32(i), 0, 0, 0}, nil)
	}
	var wg sync.WaitGroup
	// One writer goroutine, eight readers — realistic ratio for the
	// staffing co-pilot use case where ingestion is occasional and
	// queries are common.
	wg.Add(1)
	go func() {
		defer wg.Done()
		for j := 0; j < 50; j++ {
			_ = idx.Add(fmt.Sprintf("hot-%d", j), []float32{float32(j), 1, 0, 0}, nil)
		}
	}()
	for r := 0; r < 8; r++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < 50; j++ {
				_, _ = idx.Search([]float32{float32(j), 0, 0, 0}, 5)
			}
		}()
	}
	wg.Wait()
}

func TestEncodeDecode_RoundTrip(t *testing.T) {
	const n = 16
	src, _ := NewIndex(IndexParams{Name: "x", Dimension: n, Distance: DistanceCosine})
	mkVec := func(i int) []float32 {
		// Each vector is a unique unit vector along axis (i mod n) with
		// a tiny perturbation on a different axis — recall=1 is robust
		// without colliding under cosine.
		v := make([]float32, n)
		v[i%n] = 1.0
		v[(i+1)%n] = 0.001
		return v
	}
	for i := 0; i < n; i++ {
		meta := json.RawMessage(fmt.Sprintf(`{"row":%d}`, i))
		if err := src.Add(fmt.Sprintf("id-%02d", i), mkVec(i), meta); err != nil {
			t.Fatal(err)
		}
	}

	var envBuf, graphBuf bytes.Buffer
	if err := src.Encode(&envBuf, &graphBuf); err != nil {
		t.Fatalf("Encode: %v", err)
	}

	dst, err := DecodeIndex(&envBuf, &graphBuf)
	if err != nil {
		t.Fatalf("DecodeIndex: %v", err)
	}
	if dst.Len() != src.Len() {
		t.Errorf("Len: src=%d dst=%d", src.Len(), dst.Len())
	}
	if dst.Params() != src.Params() {
		t.Errorf("Params: src=%+v dst=%+v", src.Params(), dst.Params())
	}
	for i := 0; i < n; i++ {
		hits, err := dst.Search(mkVec(i), 1)
		if err != nil {
			t.Fatal(err)
		}
		want := fmt.Sprintf("id-%02d", i)
		if len(hits) == 0 || hits[0].ID != want {
			t.Errorf("Search after decode: id-%d → %v, want %s", i, hits, want)
			continue
		}
		wantMeta := fmt.Sprintf(`{"row":%d}`, i)
		if string(hits[0].Metadata) != wantMeta {
			t.Errorf("metadata after decode: got %q, want %q", hits[0].Metadata, wantMeta)
		}
	}
}

func TestDecodeIndex_VersionMismatch(t *testing.T) {
	bad := bytes.NewBufferString(`{"version":999,"params":{"name":"x","dimension":4}}`)
	_, err := DecodeIndex(bad, bytes.NewReader(nil))
	if !errors.Is(err, ErrVersionMismatch) {
		t.Errorf("expected ErrVersionMismatch, got %v", err)
	}
}

func TestRegistry_CreateGetDelete(t *testing.T) {
	r := NewRegistry()
	idx, err := r.Create(IndexParams{Name: "workers", Dimension: 4})
	if err != nil {
		t.Fatal(err)
	}
	if idx.Params().Name != "workers" {
		t.Errorf("name: got %q", idx.Params().Name)
	}
	got, err := r.Get("workers")
	if err != nil || got != idx {
		t.Errorf("Get returned different / err: %v", err)
	}
	if _, err := r.Create(IndexParams{Name: "workers", Dimension: 4}); !errors.Is(err, ErrIndexAlreadyExists) {
		t.Errorf("dup create: expected ErrIndexAlreadyExists, got %v", err)
	}
	if err := r.Delete("workers"); err != nil {
		t.Fatal(err)
	}
	if _, err := r.Get("workers"); !errors.Is(err, ErrIndexNotFound) {
		t.Errorf("Get after Delete: expected ErrIndexNotFound, got %v", err)
	}
	if err := r.Delete("workers"); !errors.Is(err, ErrIndexNotFound) {
		t.Errorf("idempotent Delete: expected ErrIndexNotFound, got %v", err)
	}
}

func TestAdd_RejectsNonFinite(t *testing.T) {
	idx, _ := NewIndex(IndexParams{Name: "x", Dimension: 3, Distance: DistanceEuclidean})
	cases := [][]float32{
		{float32(math.NaN()), 0, 0},
		{float32(math.Inf(1)), 0, 0},
		{0, float32(math.Inf(-1)), 0},
	}
	for _, vec := range cases {
		if err := idx.Add("a", vec, nil); err == nil {
			t.Errorf("expected error for non-finite vec %v", vec)
		}
	}
}

func TestAdd_RejectsZeroNormUnderCosine(t *testing.T) {
	idx, _ := NewIndex(IndexParams{Name: "x", Dimension: 3, Distance: DistanceCosine})
	if err := idx.Add("a", []float32{0, 0, 0}, nil); err == nil {
		t.Error("expected error for zero-norm under cosine")
	}
	// Same vec is OK under euclidean (origin is a valid point).
	idxE, _ := NewIndex(IndexParams{Name: "y", Dimension: 3, Distance: DistanceEuclidean})
	if err := idxE.Add("a", []float32{0, 0, 0}, nil); err != nil {
		t.Errorf("zero vec under euclidean should be valid: %v", err)
	}
}

func TestAdd_PreservesMetaOnNilReAdd(t *testing.T) {
	// Per scrum K-B1: re-adding with nil meta must NOT clear existing.
	idx, _ := NewIndex(IndexParams{Name: "x", Dimension: 2})
	_ = idx.Add("alice", []float32{1, 0}, json.RawMessage(`{"role":"warehouse"}`))
	if err := idx.Add("alice", []float32{0.9, 0.1}, nil); err != nil {
		t.Fatal(err)
	}
	_, meta, ok := idx.Lookup("alice")
	if !ok {
		t.Fatal("Lookup not found")
	}
	if string(meta) != `{"role":"warehouse"}` {
		t.Errorf("metadata cleared on nil re-add: got %q", meta)
	}
	// Explicit empty {} replaces.
	_ = idx.Add("alice", []float32{1, 0}, json.RawMessage(`{}`))
	_, meta, _ = idx.Lookup("alice")
	if string(meta) != `{}` {
		t.Errorf("explicit {} should replace: got %q", meta)
	}
}

func TestLookup_ReturnsCopy(t *testing.T) {
	// Per scrum O-W1: caller mutation must not corrupt index state.
	idx, _ := NewIndex(IndexParams{Name: "x", Dimension: 3})
	orig := []float32{1, 2, 3}
	_ = idx.Add("a", orig, nil)
	got, _, _ := idx.Lookup("a")
	got[0] = 99 // mutate the returned copy
	again, _, _ := idx.Lookup("a")
	if again[0] != 1 {
		t.Errorf("Lookup didn't copy: index now sees %v", again)
	}
}

// TestIndex_IDs locks the snapshot semantics: IDs() returns a copy
// of the metadata keyset that callers can iterate without holding
// the index lock. Underpins the merge endpoint (OPEN #1) — without
// IDs(), the merge handler can't enumerate items to drain.
func TestIndex_IDs(t *testing.T) {
	idx, err := NewIndex(IndexParams{Name: "ids_test", Dimension: 4})
	if err != nil {
		t.Fatalf("NewIndex: %v", err)
	}
	if got := idx.IDs(); len(got) != 0 {
		t.Errorf("empty index should have no IDs, got %v", got)
	}
	// Add with nil meta — the ids tracker is the canonical set, so
	// these MUST appear in IDs() even though they're not in i.meta.
	for _, id := range []string{"a", "b", "c"} {
		if err := idx.Add(id, []float32{1, 0, 0, 0}, nil); err != nil {
			t.Fatalf("Add %s: %v", id, err)
		}
	}
	got := idx.IDs()
	if len(got) != 3 {
		t.Errorf("expected 3 IDs after 3 Adds (nil meta still counts), got %d %v", len(got), got)
	}
	got[0] = "MUTATED"
	got2 := idx.IDs()
	for _, id := range got2 {
		if id == "MUTATED" {
			t.Errorf("IDs() must return a snapshot independent of internal state")
		}
	}
	// Delete updates the tracker.
	idx.Delete("a")
	if got := idx.IDs(); len(got) != 2 {
		t.Errorf("expected 2 IDs after Delete, got %d %v", len(got), got)
	}
}

func TestRegistry_Names_Sorted(t *testing.T) {
	r := NewRegistry()
	for _, n := range []string{"zoo", "alpha", "midway"} {
		_, _ = r.Create(IndexParams{Name: n, Dimension: 4})
	}
	got := r.Names()
	want := []string{"alpha", "midway", "zoo"}
	for i, w := range want {
		if got[i] != w {
			t.Errorf("Names[%d]: got %q, want %q", i, got[i], w)
		}
	}
}