Cyrene_v2/internal/runtime/memory/vector/sqlite_test.go

package vector

import (
	"context"
	"errors"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding"
)

func TestSQLite_New_NilDB_异常(t *testing.T) {
	t.Parallel()
	if _, err := NewSQLite(context.Background(), nil, 4, testLogger()); !errors.Is(err, ErrNilDB) {
		t.Fatalf("NewSQLite(nil db) err = %v, want ErrNilDB", err)
	}
}

func TestSQLite_New_InvalidDim_异常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	if _, err := NewSQLite(ctx, openTestDB(t), 0, testLogger()); !errors.Is(err, ErrInvalidDim) {
		t.Fatalf("NewSQLite(dim 0) err = %v, want ErrInvalidDim", err)
	}
}

func TestSQLite_AddSearch_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 3)

	for _, e := range []struct {
		id string
		v  Vector
	}{
		{id: "cat", v: Vector{1, 0.5, 0}},
		{id: "dog", v: Vector{0.5, 1, 0}},
		{id: "bird", v: Vector{0, 0, 1}},
	} {
		if err := s.Add(ctx, e.id, e.v); err != nil {
			t.Fatalf("Add(%s): %v", e.id, err)
		}
	}

	hits, err := s.Search(ctx, Vector{1, 0, 0}, 2)
	if err != nil {
		t.Fatalf("Search: %v", err)
	}
	if len(hits) != 2 {
		t.Fatalf("got %d hits, want 2", len(hits))
	}
	if hits[0].ID != "cat" || hits[1].ID != "dog" {
		t.Fatalf("hits = [%s %s], want [cat dog]", hits[0].ID, hits[1].ID)
	}
	if hits[0].Score <= hits[1].Score {
		t.Fatalf("top score not descending: %v <= %v", hits[0].Score, hits[1].Score)
	}
	if hits[0].Score <= 0.89 || hits[0].Score >= 0.9 {
		t.Errorf("cat score = %v, want ~0.894", hits[0].Score)
	}
}

func TestSQLite_Add_Upsert_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 3)

	if err := s.Add(ctx, "a", Vector{1, 0, 0}); err != nil {
		t.Fatalf("Add first: %v", err)
	}
	if err := s.Add(ctx, "a", Vector{0, 1, 0}); err != nil {
		t.Fatalf("Add update: %v", err)
	}

	hits, err := s.Search(ctx, Vector{0, 1, 0}, 5)
	if err != nil {
		t.Fatalf("Search: %v", err)
	}
	if len(hits) != 1 || hits[0].ID != "a" {
		t.Fatalf("hits = %+v, want single a", hits)
	}
	if hits[0].Score < 0.99 {
		t.Fatalf("a score = %v, want ~1.0 after upsert", hits[0].Score)
	}
}

func TestSQLite_Add_MissingID_异常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 4)
	if err := s.Add(ctx, "   ", Vector{1, 0, 0, 0}); !errors.Is(err, ErrMissingID) {
		t.Fatalf("Add(blank id) err = %v, want ErrMissingID", err)
	}
}

func TestSQLite_Add_DimMismatch_异常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 4)
	if err := s.Add(ctx, "x", Vector{1, 2}); !errors.Is(err, ErrDimMismatch) {
		t.Fatalf("Add(wrong dim) err = %v, want ErrDimMismatch", err)
	}
}

func TestSQLite_Search_Empty_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 4)

	hits, err := s.Search(ctx, Vector{1, 0, 0, 0}, 3)
	if err != nil {
		t.Fatalf("Search empty: %v", err)
	}
	if hits == nil || len(hits) != 0 {
		t.Fatalf("Search(empty) = %v, want non-nil empty slice", hits)
	}
}

func TestSQLite_Search_TopKNonPositive_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 4)
	if err := s.Add(ctx, "a", Vector{1, 0, 0, 0}); err != nil {
		t.Fatalf("Add: %v", err)
	}

	for _, k := range []int{0, -1} {
		hits, err := s.Search(ctx, Vector{1, 0, 0, 0}, k)
		if err != nil {
			t.Fatalf("Search(topK %d): %v", k, err)
		}
		if len(hits) != 0 {
			t.Fatalf("Search(topK %d) = %v, want empty", k, hits)
		}
	}
}

func TestSQLite_Search_DimMismatch_异常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 4)
	if _, err := s.Search(ctx, Vector{1, 0}, 3); !errors.Is(err, ErrDimMismatch) {
		t.Fatalf("Search(wrong dim) err = %v, want ErrDimMismatch", err)
	}
}

func TestSQLite_Embed_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 4)
	s.withEmbedder(&mockEmbedder{dim: 4})

	vs, err := s.Embed(ctx, []string{"hello", "world"})
	if err != nil {
		t.Fatalf("Embed: %v", err)
	}
	if len(vs) != 2 || len(vs[0]) != 4 || len(vs[1]) != 4 {
		t.Fatalf("Embed => %d vectors, want 2 x 4-dim", len(vs))
	}
}

func TestSQLite_Embed_NoEmbedder_异常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 4)
	if _, err := s.Embed(ctx, []string{"hi"}); !errors.Is(err, ErrNoEmbedder) {
		t.Fatalf("Embed(no embedder) err = %v, want ErrNoEmbedder", err)
	}
}

func TestSQLite_Embed_EmbedderError_异常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newInMemorySQLite(t, 4)
	boom := errors.New("embed boom")
	s.withEmbedder(&mockEmbedder{
		dim: 4,
		fn:  func(context.Context, []string) ([]embedding.Vector, error) { return nil, boom },
	})

	if _, err := s.Embed(ctx, []string{"hi"}); !errors.Is(err, boom) {
		t.Fatalf("Embed(driver error) err = %v, want wrap of %v", err, boom)
	}
}

func TestSQLite_Close_正常(t *testing.T) {
	t.Parallel()
	s := newInMemorySQLite(t, 4)
	if err := s.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
}