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