rony-chat-bot/internal/embed/embed_test.go

package embed

import (
	"context"
	"encoding/json"
	"math"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestNormalizeProducesUnitVectors(t *testing.T) {
	got := Normalize([]float32{3, 4})
	if math.Abs(float64(got[0])-0.6) > 1e-6 || math.Abs(float64(got[1])-0.8) > 1e-6 {
		t.Errorf("Normalize([3,4]) = %v, want [0.6 0.8]", got)
	}
	// A zero vector must survive untouched — dividing by its length would
	// put NaN into every later comparison.
	zero := []float32{0, 0, 0}
	if out := Normalize(zero); out[0] != 0 || math.IsNaN(float64(out[0])) {
		t.Errorf("Normalize(zero) = %v, want zeros", out)
	}
}

func TestSimilarity(t *testing.T) {
	a := Normalize([]float32{1, 0})
	if s := Similarity(a, a); math.Abs(s-1) > 1e-6 {
		t.Errorf("self-similarity = %v, want 1", s)
	}
	if s := Similarity(a, Normalize([]float32{0, 1})); math.Abs(s) > 1e-6 {
		t.Errorf("orthogonal similarity = %v, want 0", s)
	}
	// A vector from a different model must never outrank a real hit.
	if s := Similarity(a, []float32{1, 0, 0}); s != 0 {
		t.Errorf("mismatched dimensions scored %v, want 0", s)
	}
	if s := Similarity(nil, nil); s != 0 {
		t.Errorf("empty vectors scored %v, want 0", s)
	}
}

func TestEncodeDecodeRoundTrip(t *testing.T) {
	in := []float32{0.5, -0.25, 1e-8, 12345.75}
	out := Decode(Encode(in))
	if len(out) != len(in) {
		t.Fatalf("round trip changed length: %d → %d", len(in), len(out))
	}
	for i := range in {
		if in[i] != out[i] {
			t.Errorf("round trip [%d]: %v → %v", i, in[i], out[i])
		}
	}
	if got := Decode([]byte{1, 2, 3}); got != nil {
		t.Errorf("Decode of a truncated blob = %v, want nil", got)
	}
}

// The API returns an `index` per item rather than promising input order, so
// the client must reorder — otherwise chunk N gets chunk M's vector and every
// later search is quietly wrong.
func TestEmbedReordersByIndexAndNormalizes(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req embedRequest
		_ = json.NewDecoder(r.Body).Decode(&req)
		_ = json.NewEncoder(w).Encode(map[string]any{
			"data": []map[string]any{
				{"index": 1, "embedding": []float32{0, 5}},
				{"index": 0, "embedding": []float32{3, 4}},
			},
		})
	}))
	defer srv.Close()

	vecs, err := New(Config{BaseURL: srv.URL}).Embed(context.Background(), []string{"first", "second"})
	if err != nil {
		t.Fatal(err)
	}
	if len(vecs) != 2 {
		t.Fatalf("got %d vectors, want 2", len(vecs))
	}
	if math.Abs(float64(vecs[0][0])-0.6) > 1e-6 {
		t.Errorf("vecs[0] = %v, want the index-0 item normalized ([0.6 0.8])", vecs[0])
	}
	if math.Abs(float64(vecs[1][1])-1) > 1e-6 {
		t.Errorf("vecs[1] = %v, want the index-1 item normalized ([0 1])", vecs[1])
	}
}

func TestEmbedRejectsShortResponse(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{
			"data": []map[string]any{{"index": 0, "embedding": []float32{1, 0}}},
		})
	}))
	defer srv.Close()

	if _, err := New(Config{BaseURL: srv.URL}).Embed(context.Background(), []string{"a", "b"}); err == nil {
		t.Fatal("want an error when the endpoint returns fewer vectors than inputs")
	}
}

func TestEmbedSurfacesHTTPError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		http.Error(w, "model not loaded", http.StatusServiceUnavailable)
	}))
	defer srv.Close()

	if _, err := New(Config{BaseURL: srv.URL}).Embed(context.Background(), []string{"a"}); err == nil {
		t.Fatal("want an error on a 503")
	}
}