package search import ( "context" "path/filepath" "strings" "testing" "time" "git.wntrmute.dev/kyle/sift/internal/dcache" ) type fakeSrc struct { name string hits []Hit warn string delay time.Duration credits int } func (f fakeSrc) Name() string { return f.name } func (f fakeSrc) Search(ctx context.Context, query string, k int) SourceOut { if f.delay > 0 { select { case <-time.After(f.delay): case <-ctx.Done(): return SourceOut{Warn: f.name + ": timeout"} } } return SourceOut{Hits: f.hits, Warn: f.warn, Credits: f.credits} } type stubEmb struct { ok bool reason string } func (s stubEmb) Embed(ctx context.Context, inputs []string) ([][]float64, error) { return nil, context.Canceled } func (s stubEmb) Acquire(ctx context.Context) (bool, string) { return s.ok, s.reason } func (s stubEmb) Release() {} func TestFanInPartial(t *testing.T) { e := &Engine{ Sources: []Source{ fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Title: "ok", Snippet: "x", Source: NameSearxNG}}}, fakeSrc{name: NameBraveLLM, warn: "brave_llm: 500"}, fakeSrc{name: NameTavily, delay: 2 * time.Second}, }, Timeout: 50 * time.Millisecond, DefaultK: 8, RerankOn: false, } r := e.Run(context.Background(), "q", 8, false, "titan", dcache.ClassFleet) if len(r.Hits) != 1 || r.Hits[0].URL != "https://ok.example/a" { t.Fatalf("hits %+v", r.Hits) } joined := strings.Join(r.Warnings, " ") if !strings.Contains(joined, "500") || !strings.Contains(joined, "timeout") { t.Fatalf("warnings %v", r.Warnings) } } func TestCacheClassAndHour(t *testing.T) { db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil) if err != nil { t.Fatal(err) } defer db.Close() now := time.Date(2026, 9, 23, 10, 0, 0, 0, time.UTC) src := fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Title: "ok", Snippet: "x", Source: NameSearxNG}}} e := &Engine{Sources: []Source{src}, DefaultK: 8, Cache: db, RerankOn: false, Now: func() time.Time { return now }} r1 := e.Run(context.Background(), "Hello World", 8, false, "titan", dcache.ClassFleet) if r1.Cached { t.Fatal("first miss") } r2 := e.Run(context.Background(), "hello world", 8, false, "titan", dcache.ClassFleet) if !r2.Cached { t.Fatal("norm/hour hit") } r3 := e.Run(context.Background(), "hello world", 8, false, "gs", dcache.ClassRestricted) if r3.Cached { t.Fatal("restricted must miss fleet cache") } e.Now = func() time.Time { return now.Add(time.Hour) } r4 := e.Run(context.Background(), "hello world", 8, false, "titan", dcache.ClassFleet) if r4.Cached { t.Fatal("next hour must miss") } } func TestNoCacheOnDegraded(t *testing.T) { db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil) if err != nil { t.Fatal(err) } defer db.Close() e := &Engine{ Sources: []Source{ fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Source: NameSearxNG}}}, fakeSrc{name: NameBraveLLM, warn: "brave_llm: 500"}, }, DefaultK: 8, Cache: db, RerankOn: false, } e.Run(context.Background(), "q", 8, false, "titan", dcache.ClassFleet) r := e.Run(context.Background(), "q", 8, false, "titan", dcache.ClassFleet) if r.Cached { t.Fatal("degraded cached") } } func TestRerankFallbackStill200(t *testing.T) { e := &Engine{ Sources: []Source{fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Source: NameSearxNG}}}}, DefaultK: 8, RerankOn: true, Embedder: stubEmb{ok: false, reason: "queue_full"}, } r := e.Run(context.Background(), "q", 8, true, "titan", dcache.ClassFleet) if len(r.Hits) != 1 { t.Fatal(r.Hits) } if strings.Contains(strings.Join(r.Warnings, " "), "rerank fell back") { t.Fatalf("RRF must not depend on embed: %v", r.Warnings) } } func TestRRFInterleaveAndNoRerank(t *testing.T) { searx := []Hit{ {URL: "https://a.example/s", Title: "s1", Source: NameSearxNG}, {URL: "https://b.example/s", Title: "s2", Source: NameSearxNG}, } tav := []Hit{ {URL: "https://tavily.example/1", Title: "t1", Source: NameTavily}, {URL: "https://a.example/s", Title: "t-dup", Snippet: "longer tavily snippet here", Source: NameTavily}, } merged := rrfMerge([][]Hit{searx, tav}) if len(merged) != 3 { t.Fatalf("%+v", merged) } if merged[0].URL != "https://a.example/s" { t.Fatalf("shared URL should rank first, got %+v", merged[0]) } e := &Engine{ Sources: []Source{ fakeSrc{name: NameSearxNG, hits: searx}, fakeSrc{name: NameTavily, hits: tav}, }, DefaultK: 8, RerankOn: true, NoRerank: map[string]bool{NameTavily: true}, } r := e.Run(context.Background(), "q", 8, true, "titan", dcache.ClassFleet) if r.Hits[0].URL != "https://tavily.example/1" { t.Fatalf("no_rerank tavily should lead: %+v", r.Hits) } } func TestRRFKeepsOriginalURL(t *testing.T) { orig := "https://ex.com/docs/" merged := rrfMerge([][]Hit{ {{URL: orig, Title: "slash", Source: NameSearxNG}}, {{URL: "https://ex.com/docs", Title: "stripped", Snippet: "longer", Source: NameTavily}}, }) if len(merged) != 1 { t.Fatalf("%+v", merged) } if merged[0].URL != orig { t.Fatalf("want original URL %q got %q", orig, merged[0].URL) } if CanonicalURL(orig) != "https://ex.com/docs" { t.Fatalf("canonical %q", CanonicalURL(orig)) } } func TestDailyCapSkipsTavily(t *testing.T) { db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil) if err != nil { t.Fatal(err) } defer db.Close() var n int tav := fakeSrc{name: NameTavily, hits: []Hit{{URL: "https://t.example/a", Source: NameTavily}}, credits: 1} e := &Engine{Sources: []Source{tav}, DefaultK: 8, Cache: db, TavilyCap: 1, RerankOn: false} e.Run(context.Background(), "q1", 8, false, "titan", dcache.ClassFleet) r := e.Run(context.Background(), "q2", 8, false, "titan", dcache.ClassFleet) for _, w := range r.Warnings { if strings.Contains(w, "daily cap") { n++ } } if n != 1 { t.Fatalf("warnings %v", r.Warnings) } } func TestRestrictedSkipsTavilyCap(t *testing.T) { db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil) if err != nil { t.Fatal(err) } defer db.Close() tav := fakeSrc{name: NameTavily, hits: []Hit{{URL: "https://t.example/a", Source: NameTavily}}, credits: 1} e := &Engine{Sources: []Source{tav}, DefaultK: 8, Cache: db, TavilyCap: 1, RerankOn: false} e.Run(context.Background(), "q1", 8, false, "titan", dcache.ClassFleet) day := e.now().UTC().Format("2006-01-02") if db.GetSearchCredits(NameTavily, day) != 1 { t.Fatalf("fleet row %d", db.GetSearchCredits(NameTavily, day)) } r := e.Run(context.Background(), "q-gs", 8, false, "straylight-gs", dcache.ClassRestricted) if len(r.Hits) != 1 || r.Hits[0].URL != "https://t.example/a" { t.Fatalf("restricted should still get Tavily: %+v %v", r.Hits, r.Warnings) } for _, w := range r.Warnings { if strings.Contains(w, "daily cap") { t.Fatalf("restricted capped: %v", r.Warnings) } } if db.GetSearchCredits(NameTavily, day) != 1 { t.Fatalf("restricted must not raise fleet row, got %d", db.GetSearchCredits(NameTavily, day)) } if db.GetSearchCredits(NameTavily+":restricted", day) != 1 { t.Fatalf("restricted row %d", db.GetSearchCredits(NameTavily+":restricted", day)) } fleet := e.Run(context.Background(), "q2", 8, false, "titan", dcache.ClassFleet) n := 0 for _, w := range fleet.Warnings { if strings.Contains(w, "daily cap") { n++ } } if n != 1 { t.Fatalf("fleet still capped: %v", fleet.Warnings) } } func TestDedupeLongerSnippet(t *testing.T) { in := []Hit{ {URL: "https://Ex.com/a#x", Snippet: "aa", Source: NameSearxNG}, {URL: "https://ex.com/a", Snippet: "aabbcc", Source: NameTavily}, } out := Dedupe(in) if len(out) != 1 || out[0].Source != NameSearxNG || out[0].Snippet != "aabbcc" { t.Fatalf("%+v", out) } if out[0].URL != "https://Ex.com/a#x" { t.Fatalf("dedupe must keep first original URL, got %q", out[0].URL) } }