package embedding import ( "context" "encoding/json" "errors" "io" "net/http" "net/http/httptest" "strings" "testing" "git.yeij.top/AskaEth/Cyrene_v2/internal/config" ) // roundTripFunc adapts a function to http.RoundTripper, letting tests inject a // mock transport without starting a network server. type roundTripFunc func(*http.Request) (*http.Response, error) func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } func TestAPIEmbedder_EmbedsFromServer(t *testing.T) { t.Parallel() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { t.Errorf("method = %s, want POST", r.Method) } if got := r.Header.Get("Content-Type"); got != "application/json" { t.Errorf("Content-Type = %q, want application/json", got) } if got := r.Header.Get("Authorization"); got != "Bearer test-key" { t.Errorf("Authorization = %q, want Bearer test-key", got) } var req openAICompatibleRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { t.Errorf("decode request: %v", err) w.WriteHeader(http.StatusBadRequest) return } if req.Model != "text-embedding-3-small" { t.Errorf("request model = %q, want text-embedding-3-small", req.Model) } vec := make([]float32, 4) for i := range vec { vec[i] = float32(i) + 0.1 } resp := openAICompatibleResponse{Data: make([]openAICompatibleDatum, 0, len(req.Input))} for i := range req.Input { resp.Data = append(resp.Data, openAICompatibleDatum{Index: i, Embedding: vec}) } _ = json.NewEncoder(w).Encode(resp) })) defer srv.Close() cfg := config.EmbeddingConfig{ APIURL: srv.URL + "/v1/embeddings", APIKey: "test-key", Model: "text-embedding-3-small", } emb := NewAPI(cfg) vectors, err := emb.Embed(context.Background(), []string{"hello", "world"}) if err != nil { t.Fatalf("Embed() error: %v", err) } if len(vectors) != 2 { t.Fatalf("got %d vectors, want 2", len(vectors)) } for i, v := range vectors { if len(v) != 4 { t.Errorf("vector %d length = %d, want 4", i, len(v)) } for j, x := range v { if want := float32(j) + 0.1; x != want { t.Errorf("vector %d[%d] = %v, want %v", i, j, x, want) } } } if emb.Dim() != 4 { t.Errorf("Dim() = %d, want 4", emb.Dim()) } } func TestAPIEmbedder_RoundTripperInjection(t *testing.T) { t.Parallel() var seenPath string rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { seenPath = r.URL.Path body := `{"data":[{"index":0,"embedding":[0.1,0.2,0.3]}]}` return &http.Response{ StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body)), Header: make(http.Header), }, nil }) emb := NewAPI( config.EmbeddingConfig{APIURL: "http://example.com/v1/embeddings"}, WithRoundTripper(rt), ) vectors, err := emb.Embed(context.Background(), []string{"hi"}) if err != nil { t.Fatalf("Embed() error: %v", err) } if seenPath != "/v1/embeddings" { t.Errorf("request path = %q, want /v1/embeddings", seenPath) } if len(vectors) != 1 || len(vectors[0]) != 3 { t.Fatalf("vectors = %v, want 1 vector of length 3", vectors) } } func TestAPIEmbedder_Non200Errors(t *testing.T) { t.Parallel() rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { return &http.Response{ StatusCode: http.StatusInternalServerError, Body: io.NopCloser(strings.NewReader("server exploded")), Header: make(http.Header), }, nil }) emb := NewAPI( config.EmbeddingConfig{APIURL: "http://example.com/v1/embeddings"}, WithRoundTripper(rt), ) if _, err := emb.Embed(context.Background(), []string{"hi"}); err == nil { t.Fatal("expected error on 500, got nil") } } func TestAPIEmbedder_MismatchedCountErrors(t *testing.T) { t.Parallel() rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { body := `{"data":[{"index":0,"embedding":[0.1,0.2]}]}` return &http.Response{ StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body)), Header: make(http.Header), }, nil }) emb := NewAPI( config.EmbeddingConfig{APIURL: "http://example.com/v1/embeddings"}, WithRoundTripper(rt), ) // 2 texts sent, only 1 returned -> error. if _, err := emb.Embed(context.Background(), []string{"a", "b"}); err == nil { t.Fatal("expected error on count mismatch, got nil") } } func TestAPIEmbedder_BlankURLErrors(t *testing.T) { t.Parallel() emb := NewAPI(config.EmbeddingConfig{}) _, err := emb.Embed(context.Background(), []string{"hello"}) if err == nil { t.Fatal("expected error on blank URL, got nil") } if !errors.Is(err, ErrNoAPIURL) { t.Errorf("error = %v, want ErrNoAPIURL", err) } } func TestAPIEmbedder_EmptyInput(t *testing.T) { t.Parallel() emb := NewAPI(config.EmbeddingConfig{APIURL: "http://example.invalid"}) vectors, err := emb.Embed(context.Background(), nil) if err != nil { t.Fatalf("Embed(nil) error: %v", err) } if len(vectors) != 0 { t.Errorf("got %d vectors, want 0", len(vectors)) } } func TestAPIEmbedder_CanceledContext(t *testing.T) { t.Parallel() emb := NewAPI(config.EmbeddingConfig{APIURL: "http://example.invalid"}) ctx, cancel := context.WithCancel(context.Background()) cancel() if _, err := emb.Embed(ctx, []string{"hello"}); err == nil { t.Fatal("expected error on canceled context, got nil") } }