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) } }) } }