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 返回一个不写任何地方的 logger，保持测试输出干净。
func testLogger() *slog.Logger {
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

// openTestDB 打开一个隔离的内存 SQLite 库。每次调用都拥有私有 DB，
// 因此并行测试天生相互隔离。返回的 DB 由
// 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 在全新的内存 DB 上构建维度为 dim 的 SQLiteIndex。
// embedder 不接入，由测试自行控制。
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 是测试用的可脚本化 embedding.Embedder。未设置 fn 时它
// 返回 one-hot 向量，每段输入文本一个，测试因此可预测结果。
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 接入了 embedder，因此返回的索引上 Embed 可用。
	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()
	// 指向已关闭端口的环回 URL 会快速失败，证明 New 选择了 PG 路径
	// 而不是静默回退到 SQLite。端口 1（tcpmux）在
	// localhost 上不监听，因此连接立即被拒。
	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)
			}
		})
	}
}