Cyrene_v2/internal/runtime/embedding/api_test.go

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