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

package vector

import (
	"context"
	"database/sql"
	"errors"
	"io"
	"log/slog"
	"testing"

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

// testLogger returns a logger that writes nowhere, keeping test output clean.
func testLogger() *slog.Logger {
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

// openTestDB opens an isolated in-memory SQLite DB. Each call owns a private DB,
// so parallel tests are isolated by construction. The returned DB is closed by
// t.Cleanup.
func openTestDB(t *testing.T) *sql.DB {
	t.Helper()
	ctx := context.Background()
	d, err := db.Open(ctx, ":memory:", testLogger())
	if err != nil {
		t.Fatalf("open in-memory db: %v", err)
	}
	t.Cleanup(func() { _ = d.Close() })
	return d
}

// newInMemorySQLite builds a SQLiteIndex on a fresh in-memory DB with dim
// dimensions. The embedder is left unwired so tests control it.
func newInMemorySQLite(t *testing.T, dim int) *SQLiteIndex {
	t.Helper()
	ctx := context.Background()
	idx, err := NewSQLite(ctx, openTestDB(t), dim, testLogger())
	if err != nil {
		t.Fatalf("NewSQLite: %v", err)
	}
	return idx
}

// mockEmbedder is a scriptable embedding.Embedder for tests. With no fn it
// returns one-hot vectors, one per input text, so tests can predict results.
type mockEmbedder struct {
	dim int
	fn  func(ctx context.Context, texts []string) ([]embedding.Vector, error)
}

func (m *mockEmbedder) Dim() int { return m.dim }

func (m *mockEmbedder) Embed(ctx context.Context, texts []string) ([]embedding.Vector, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if m.fn != nil {
		return m.fn(ctx, texts)
	}
	out := make([]embedding.Vector, len(texts))
	for i := range texts {
		if m.dim <= 0 {
			out[i] = embedding.Vector{}
			continue
		}
		v := make(embedding.Vector, m.dim)
		v[i%m.dim] = 1
		out[i] = v
	}
	return out, nil
}

func TestNew_NilEmbedder_异常(t *testing.T) {
	t.Parallel()
	if _, err := New(config.DBConfig{}, nil, testLogger()); !errors.Is(err, ErrNoEmbedder) {
		t.Fatalf("New(nil embedder) err = %v, want ErrNoEmbedder", err)
	}
}

func TestNew_InvalidDim_异常(t *testing.T) {
	t.Parallel()
	emb := &mockEmbedder{dim: 0}
	if _, err := New(config.DBConfig{SQLitePath: ":memory:"}, emb, testLogger()); !errors.Is(err, ErrInvalidDim) {
		t.Fatalf("New(dim 0) err = %v, want ErrInvalidDim", err)
	}
}

func TestNew_SelectsSQLite_WhenPGURLBlank(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	idx, err := New(config.DBConfig{SQLitePath: ":memory:"}, &mockEmbedder{dim: 4}, testLogger())
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	if _, ok := idx.(*SQLiteIndex); !ok {
		t.Fatalf("New type = %T, want *SQLiteIndex", idx)
	}
	// New wires the embedder, so Embed works on the returned index.
	vs, err := idx.Embed(ctx, []string{"hello"})
	if err != nil {
		t.Fatalf("Embed: %v", err)
	}
	if len(vs) != 1 || len(vs[0]) != 4 {
		t.Fatalf("Embed => %v, want one 4-dim vector", vs)
	}
	if err := idx.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
}

func TestNew_PGURLWhitespace_SelectsSQLite(t *testing.T) {
	t.Parallel()
	idx, err := New(config.DBConfig{SQLitePath: ":memory:", PGURL: "  \t "}, &mockEmbedder{dim: 4}, testLogger())
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	if _, ok := idx.(*SQLiteIndex); !ok {
		t.Fatalf("New type = %T, want *SQLiteIndex for whitespace PGURL", idx)
	}
	if err := idx.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
}

func TestNew_PGSelection_AttemptsPGWhenURLSet(t *testing.T) {
	t.Parallel()
	// A loopback URL on a closed port fails fast, proving New chose the PG path
	// rather than silently falling back to SQLite. Port 1 (tcpmux) does not
	// listen on localhost, so the connection is refused immediately.
	idx, err := New(config.DBConfig{
		SQLitePath: ":memory:",
		PGURL:      "postgres://127.0.0.1:1/mem?sslmode=disable&connect_timeout=1",
	}, &mockEmbedder{dim: 4}, testLogger())
	if err == nil {
		if idx != nil {
			_ = idx.Close()
		}
		t.Fatal("New with unreachable PGURL: expected error, got nil")
	}
	if idx != nil {
		t.Fatalf("New returned non-nil index %T on error", idx)
	}
}

func TestUsePG(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name string
		url  string
		want bool
	}{
		{name: "empty", url: "", want: false},
		{name: "whitespace", url: "  \t ", want: false},
		{name: "set", url: "postgres://user:pass@host:5432/db", want: true},
	}
	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			if got := usePG(tt.url); got != tt.want {
				t.Fatalf("usePG(%q) = %v, want %v", tt.url, got, tt.want)
			}
		})
	}
}