sift/internal/search/sources_test.go

package search

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/search/dial"
)

func TestTavilyOutboundAndDropsAnswer(t *testing.T) {
	var gotBody []byte
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			t.Errorf("method %s", r.Method)
		}
		if !strings.HasPrefix(r.Header.Get("Authorization"), "Bearer ") {
			t.Error("missing bearer")
		}
		gotBody, _ = io.ReadAll(r.Body)
		_, _ = w.Write([]byte(`{"answer":"SECRET","results":[{"url":"https://ex.com/a","title":"A","content":"snip","score":0.9}],"usage":{"credits":1}}`))
	}))
	defer ts.Close()
	c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: time.Second})
	if err != nil {
		t.Fatal(err)
	}
	out := (&Tavily{Client: c, Key: "test-key"}).Search(context.Background(), "q", 5, Filter{})
	if out.Warn != "" {
		t.Fatal(out.Warn)
	}
	var req map[string]any
	if err := json.Unmarshal(gotBody, &req); err != nil {
		t.Fatal(err)
	}
	if req["include_answer"] != false || req["include_raw_content"] != false || req["include_images"] != false || req["search_depth"] != "basic" {
		t.Fatalf("%+v", req)
	}
	if req["auto_parameters"] != false {
		t.Fatalf("auto_parameters %+v", req["auto_parameters"])
	}
	if strings.Contains(string(gotBody), "extract") || strings.Contains(string(gotBody), "crawl") {
		t.Fatal(string(gotBody))
	}
	if len(out.Hits) != 1 || out.Hits[0].URL != "https://ex.com/a" {
		t.Fatalf("%+v", out.Hits)
	}
	for _, h := range out.Hits {
		if strings.Contains(h.Snippet, "SECRET") || strings.Contains(h.Title, "SECRET") {
			t.Fatal("answer leaked into hit")
		}
	}
}

func TestSearxNGMapsAndSorts(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/search" || r.URL.Query().Get("format") != "json" {
			t.Errorf("path %s", r.URL.String())
		}
		_, _ = w.Write([]byte(`{"results":[
			{"url":"https://b.example/","title":"b","content":"bb","score":0.1},
			{"url":"https://a.example/","title":"a","content":"aa","score":0.9}
		]}`))
	}))
	defer ts.Close()
	c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: time.Second})
	if err != nil {
		t.Fatal(err)
	}
	out := (&SearxNG{Client: c}).Search(context.Background(), "q", 8, Filter{})
	if out.Warn != "" || len(out.Hits) != 2 || out.Hits[0].Title != "a" {
		t.Fatalf("%+v", out)
	}
}

func TestBraveMapsSources(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("X-Subscription-Token") == "" {
			t.Error("token")
		}
		_, _ = w.Write([]byte(`{"grounding":"ignore me","sources":[{"url":"https://s.example/x","title":"S","snippet":"hello"}]}`))
	}))
	defer ts.Close()
	c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: time.Second})
	if err != nil {
		t.Fatal(err)
	}
	out := (&BraveLLM{Client: c, Key: "k"}).Search(context.Background(), "q", 5, Filter{})
	if len(out.Hits) != 1 || out.Hits[0].URL != "https://s.example/x" {
		t.Fatalf("%+v", out)
	}
	if strings.Contains(out.Hits[0].Snippet, "ignore") {
		t.Fatal("grounding used as snippet")
	}
}

func TestBraveMapsGroundingGeneric(t *testing.T) {
	raw, err := os.ReadFile("../../tests/testdata/brave_llm_context.json")
	if err != nil {
		t.Fatal(err)
	}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write(raw)
	}))
	defer ts.Close()
	c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: time.Second})
	if err != nil {
		t.Fatal(err)
	}
	out := (&BraveLLM{Client: c, Key: "k"}).Search(context.Background(), "q", 5, Filter{})
	if len(out.Hits) != 3 {
		t.Fatalf("hits %d %+v", len(out.Hits), out.Hits)
	}
	if out.Hits[0].URL != "https://docs.tavily.com/documentation/api-reference/endpoint/search" {
		t.Fatalf("first %+v", out.Hits[0])
	}
	if !strings.Contains(out.Hits[0].Snippet, "max_results") {
		t.Fatalf("snippet %q", out.Hits[0].Snippet)
	}
	var urls []string
	for _, h := range out.Hits {
		urls = append(urls, h.URL)
	}
	joined := strings.Join(urls, " ")
	if !strings.Contains(joined, "orphan.example") {
		t.Fatalf("sources map missed: %v", urls)
	}
}

func TestTavilyBadJSONIsError(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`not-json`))
	}))
	defer ts.Close()
	c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: time.Second})
	if err != nil {
		t.Fatal(err)
	}
	out := (&Tavily{Client: c, Key: "k"}).Search(context.Background(), "q", 5, Filter{})
	if out.Warn != "tavily: error" {
		t.Fatalf("warn %q", out.Warn)
	}
	if len(out.Hits) != 0 || out.Credits != 0 {
		t.Fatalf("%+v", out)
	}
}