fathom/internal/insights/local/local_test.go

package local

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

// TestCleanLabel is where a small model's habits get absorbed.
//
// Local models narrate, apologise, emit reasoning blocks, and wrap answers in
// quotes. A label is shown to the reader as a finding's name, so anything that
// does not look like a label must be rejected in favour of the heuristic title.
// Rejecting is always safe; accepting rubbish is not.
func TestCleanLabel(t *testing.T) {
	cases := map[string]string{
		"Flaky test triage":                           "Flaky test triage",
		"  Flaky test triage.  ":                      "Flaky test triage",
		`"Flaky test triage"`:                         "Flaky test triage",
		"**Flaky test triage**":                       "Flaky test triage",
		"- Flaky test triage":                         "Flaky test triage",
		"Common theme: Flaky test triage":             "Flaky test triage",
		"Sure! Here you go:\nFlaky test triage":       "Flaky test triage",
		"<think>hmm they all...</think>\nFlaky tests": "Flaky tests",

		// Prose means the instruction was ignored; the caller keeps its own title.
		"Sure, I'd be happy to help you with that! Looking at these prompts, it seems like they are all about tests": "",
		"": "",
	}
	for in, want := range cases {
		if got := cleanLabel(in); got != want {
			t.Errorf("cleanLabel(%q) = %q, want %q", in, got, want)
		}
	}
}

func TestPickModels(t *testing.T) {
	embed, chat := pickModels([]string{"llama3.2:3b", "nomic-embed-text:latest", "qwen2.5:7b"})
	if embed != "nomic-embed-text:latest" {
		t.Errorf("embed = %q, want the embedding model", embed)
	}
	// Small models are preferred: the job is writing a six-word phrase.
	if chat != "llama3.2:3b" {
		t.Errorf("chat = %q, want the smaller chat model", chat)
	}

	// An embedding model must never be chosen to chat with.
	embed, chat = pickModels([]string{"mxbai-embed-large"})
	if embed != "mxbai-embed-large" {
		t.Errorf("embed = %q", embed)
	}
	if chat != "" {
		t.Errorf("chat = %q, want empty when only an embedding model is present", chat)
	}

	// An unrecognised name is still usable for chat.
	_, chat = pickModels([]string{"some-private-model:v3"})
	if chat != "some-private-model:v3" {
		t.Errorf("chat = %q, want the unrecognised model as a fallback", chat)
	}
}

// fakeOllama serves enough of the API to exercise discovery and both calls.
func fakeOllama(t *testing.T, models []string) *httptest.Server {
	t.Helper()
	mux := http.NewServeMux()
	mux.HandleFunc("/api/tags", func(w http.ResponseWriter, r *http.Request) {
		out := map[string]any{"models": []map[string]string{}}
		var list []map[string]string
		for _, m := range models {
			list = append(list, map[string]string{"name": m})
		}
		out["models"] = list
		json.NewEncoder(w).Encode(out)
	})
	mux.HandleFunc("/api/embed", func(w http.ResponseWriter, r *http.Request) {
		var req struct {
			Input []string `json:"input"`
		}
		json.NewDecoder(r.Body).Decode(&req)
		vecs := make([][]float32, len(req.Input))
		for i := range vecs {
			vecs[i] = []float32{1, float32(i), 0}
		}
		json.NewEncoder(w).Encode(map[string]any{"embeddings": vecs})
	})
	mux.HandleFunc("/api/generate", func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(map[string]any{"response": "Flaky test triage"})
	})
	return httptest.NewServer(mux)
}

func TestDiscoverAndUseOllama(t *testing.T) {
	srv := fakeOllama(t, []string{"llama3.2:3b", "nomic-embed-text"})
	defer srv.Close()
	t.Setenv("OLLAMA_HOST", srv.URL)

	p := Discover(context.Background())
	if p == nil {
		t.Fatal("did not discover the server named by OLLAMA_HOST")
	}
	if p.Kind != KindOllama {
		t.Errorf("kind = %q, want %q", p.Kind, KindOllama)
	}
	if !p.CanEmbed() || !p.CanLabel() {
		t.Fatalf("capabilities: embed=%v label=%v", p.CanEmbed(), p.CanLabel())
	}

	vecs, err := p.Embed(context.Background(), []string{"a", "b"})
	if err != nil {
		t.Fatal(err)
	}
	if len(vecs) != 2 || len(vecs[0]) != 3 {
		t.Errorf("got %d vectors of width %d", len(vecs), len(vecs[0]))
	}

	label, err := p.Label(context.Background(), []string{"fix the flaky test", "fix the other flaky test"})
	if err != nil {
		t.Fatal(err)
	}
	if label != "Flaky test triage" {
		t.Errorf("label = %q", label)
	}
}

func TestDiscoverOpenAICompatible(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(map[string]any{
			"data": []map[string]string{{"id": "qwen2.5-3b-instruct"}, {"id": "text-embedding-bge-small"}},
		})
	})
	mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(map[string]any{
			"choices": []map[string]any{{"message": map[string]string{"content": "Pagination work"}}},
		})
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()
	t.Setenv("FATHOM_LOCAL_MODEL_URL", srv.URL)

	p := Discover(context.Background())
	if p == nil {
		t.Fatal("did not discover the OpenAI-compatible server")
	}
	if p.Kind != KindOpenAI {
		t.Errorf("kind = %q, want %q", p.Kind, KindOpenAI)
	}
	if p.EmbedModel != "text-embedding-bge-small" {
		t.Errorf("embed model = %q", p.EmbedModel)
	}

	label, err := p.Label(context.Background(), []string{"add pagination", "add paging"})
	if err != nil {
		t.Fatal(err)
	}
	if label != "Pagination work" {
		t.Errorf("label = %q", label)
	}
}

// TestDiscoverFindsNothing is the ordinary case on a machine with no model
// server. It must return nil quickly and without an error.
func TestDiscoverFindsNothing(t *testing.T) {
	// A server that exists but speaks neither protocol.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.NotFound(w, r)
	}))
	defer srv.Close()
	t.Setenv("OLLAMA_HOST", srv.URL)
	t.Setenv("FATHOM_LOCAL_MODEL_URL", srv.URL)

	if p := Discover(context.Background()); p != nil {
		t.Errorf("discovered %+v from a server that speaks neither protocol", p)
	}
}

func TestEmbedWithoutModelIsAnError(t *testing.T) {
	p := &Provider{BaseURL: "http://localhost:1", Kind: KindOllama}
	if _, err := p.Embed(context.Background(), []string{"x"}); err == nil {
		t.Error("embedding with no embed model should error so the caller falls back")
	}
	if _, err := p.Label(context.Background(), []string{"x"}); err == nil {
		t.Error("labelling with no chat model should error so the caller keeps its title")
	}
}

// TestCandidatesPrefersExplicitSettings: an operator who names an endpoint
// should not have it lose to a default port.
func TestCandidatesPrefersExplicitSettings(t *testing.T) {
	t.Setenv("OLLAMA_HOST", "myhost:9999")
	got := candidates()
	if len(got) == 0 || got[0] != "http://myhost:9999" {
		t.Errorf("first candidate = %v, want the OLLAMA_HOST value with a scheme added", got)
	}
	if !strings.Contains(strings.Join(got, ","), "11434") {
		t.Error("the default Ollama port should still be probed as a fallback")
	}
}