package reranker_test import ( "context" "encoding/json" "io" "net/http" "net/http/httptest" "strings" "testing" "github.com/mathiasbq/hyperguild/ingestion/internal/reranker" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // fakeOllama responds to /api/generate based on a per-document // {needle → answer} map: if the prompt contains the needle, returns // the mapped answer. type fakeOllama struct { t *testing.T answers map[string]string // needle → "yes" or "no" calls int lastBody map[string]any } func (f *fakeOllama) handler() http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(f.t, http.MethodPost, r.Method) require.Equal(f.t, "/api/generate", r.URL.Path) body, err := io.ReadAll(r.Body) require.NoError(f.t, err) var p map[string]any require.NoError(f.t, json.Unmarshal(body, &p)) f.calls++ f.lastBody = p prompt := p["prompt"].(string) answer := "no" for needle, a := range f.answers { if strings.Contains(prompt, needle) { answer = a break } } w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(map[string]any{ "model": p["model"], "response": answer, "done": true, }) }) } func TestNew_EmptyURLReturnsNil(t *testing.T) { assert.Nil(t, reranker.New("", "model")) } func TestScore_YesAndNoOrdered(t *testing.T) { f := &fakeOllama{t: t, answers: map[string]string{ "alpha doc": "yes", "beta doc": "no", "gamma doc": "yes", }} srv := httptest.NewServer(f.handler()) defer srv.Close() c := reranker.New(srv.URL, "test-model") require.NotNil(t, c) scores, err := c.Score(context.Background(), "what is alpha", []string{"alpha doc body", "beta doc body", "gamma doc body"}) require.NoError(t, err) require.Len(t, scores, 3) assert.Equal(t, 1.0, scores[0]) assert.Equal(t, 0.0, scores[1]) assert.Equal(t, 1.0, scores[2]) assert.Equal(t, 3, f.calls) } func TestScore_SendsCorrectShape(t *testing.T) { f := &fakeOllama{t: t, answers: map[string]string{"hello": "yes"}} srv := httptest.NewServer(f.handler()) defer srv.Close() c := reranker.New(srv.URL, "qwen3-rerank") _, err := c.Score(context.Background(), "greeting", []string{"hello world"}) require.NoError(t, err) assert.Equal(t, "qwen3-rerank", f.lastBody["model"]) prompt := f.lastBody["prompt"].(string) assert.Contains(t, prompt, "greeting") assert.Contains(t, prompt, "hello world") assert.Contains(t, prompt, `"yes" or "no"`) } func TestScore_HandlesAmbiguousResponse(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _ = json.NewEncoder(w).Encode(map[string]any{"response": "maybe — unclear", "done": true}) })) defer srv.Close() c := reranker.New(srv.URL, "m") scores, err := c.Score(context.Background(), "q", []string{"d"}) require.NoError(t, err) // Anything that does not start with "yes" (case-insensitive, after // whitespace/think trim) is treated as "no" = 0. assert.Equal(t, []float64{0}, scores) } func TestScore_EmptyDocsReturnsEmpty(t *testing.T) { c := reranker.New("http://127.0.0.1:1", "m") scores, err := c.Score(context.Background(), "q", nil) require.NoError(t, err) assert.Empty(t, scores) } func TestScore_UpstreamErrorPropagates(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusInternalServerError) })) defer srv.Close() c := reranker.New(srv.URL, "m") _, err := c.Score(context.Background(), "q", []string{"d"}) require.Error(t, err) }