108 lines
4.2 KiB
Go
108 lines
4.2 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
// The results of DuckDuckGo's light page: titles, real addresses (not its
|
|
// redirects), snippets without tags.
|
|
func TestParseDDGLite(t *testing.T) {
|
|
b, err := os.ReadFile("testdata/ddg_lite.html")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
rs := parseDDG(string(b))
|
|
if len(rs) < 5 {
|
|
t.Fatalf("results: %d", len(rs))
|
|
}
|
|
r := rs[0]
|
|
if r.URL != "https://go.dev/doc/database/cancel-operations" || !strings.Contains(r.Title, "Canceling in-progress operations") ||
|
|
!strings.Contains(r.Snippet, "Learn how to use context.Context") || strings.Contains(r.Snippet, "<b>") {
|
|
t.Fatalf("first: %+v", r)
|
|
}
|
|
}
|
|
|
|
// Each search engine of the settings, and DuckDuckGo when the chosen one
|
|
// fails.
|
|
func TestSearchProviders(t *testing.T) {
|
|
var gotKey, gotQ string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch {
|
|
case r.URL.Path == "/search" && r.URL.Query().Get("format") == "json": // SearXNG
|
|
gotQ = r.URL.Query().Get("q")
|
|
fmt.Fprint(w, `{"results":[{"title":"SX","url":"https://a.example/","content":"from searx"}]}`)
|
|
case r.URL.Path == "/res/v1/web/search": // Brave
|
|
gotKey, gotQ = r.Header.Get("X-Subscription-Token"), r.URL.Query().Get("q")
|
|
fmt.Fprint(w, `{"web":{"results":[{"title":"BR","url":"https://b.example/","description":"from <strong>brave</strong>"}]}}`)
|
|
case r.URL.Path == "/search" && r.Method == http.MethodPost: // Tavily
|
|
var in map[string]any
|
|
_ = json.NewDecoder(r.Body).Decode(&in)
|
|
gotKey, gotQ = r.Header.Get("Authorization"), in["query"].(string)
|
|
fmt.Fprint(w, `{"results":[{"title":"TV","url":"https://t.example/","content":"from tavily"}]}`)
|
|
default:
|
|
http.Error(w, "no", 500)
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
braveURL, tavilyURL = srv.URL+"/res/v1/web/search", srv.URL+"/search"
|
|
ctx := context.Background()
|
|
|
|
rs, used, err := searchWith(ctx, searchConf{Engine: "searxng", URL: srv.URL}, "q1", 6)
|
|
if err != nil || len(rs) != 1 || rs[0].Title != "SX" || gotQ != "q1" || used != "SearXNG" {
|
|
t.Fatalf("searxng: %+v %v", rs, err)
|
|
}
|
|
rs, _, err = searchWith(ctx, searchConf{Engine: "brave", Key: "bk"}, "q2", 6)
|
|
if err != nil || rs[0].Snippet != "from brave" || gotKey != "bk" || gotQ != "q2" {
|
|
t.Fatalf("brave: %+v %v %q", rs, err, gotKey)
|
|
}
|
|
rs, used, err = searchWith(ctx, searchConf{Engine: "tavily", Key: "tk"}, "q3", 6)
|
|
if err != nil || rs[0].Title != "TV" || gotKey != "Bearer tk" || gotQ != "q3" || used != "Tavily" {
|
|
t.Fatalf("tavily: %+v %v %q", rs, err, gotKey)
|
|
}
|
|
// The chosen one fails: DuckDuckGo instead (stood in for here).
|
|
ddgSearch = func(ctx context.Context, q string, n int) ([]searchResult, error) {
|
|
return []searchResult{{Title: "DDG"}}, nil
|
|
}
|
|
rs, used, err = searchWith(ctx, searchConf{Engine: "searxng", URL: srv.URL + "/broken"}, "q4", 6)
|
|
if err != nil || rs[0].Title != "DDG" || !strings.Contains(used, "DuckDuckGo") || !strings.Contains(used, "SearXNG не ответил") {
|
|
t.Fatalf("fallback: %+v %v", rs, err)
|
|
}
|
|
}
|
|
|
|
// The chosen engine's limit spent (429 and the like): DuckDuckGo answers.
|
|
func TestSearchLimitFallsBack(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Error(w, `{"error":"rate limit"}`, http.StatusTooManyRequests)
|
|
}))
|
|
defer srv.Close()
|
|
old := ddgSearch
|
|
defer func() { ddgSearch = old }()
|
|
ddgSearch = func(ctx context.Context, q string, n int) ([]searchResult, error) {
|
|
return []searchResult{{Title: "DDG " + q}}, nil
|
|
}
|
|
braveURL, tavilyURL = srv.URL, srv.URL
|
|
for _, e := range []string{"brave", "tavily", "searxng"} {
|
|
rs, _, err := searchWith(context.Background(), searchConf{Engine: e, Key: "k", URL: srv.URL}, "x", 6)
|
|
if err != nil || len(rs) != 1 || rs[0].Title != "DDG x" {
|
|
t.Fatalf("%s: %+v %v", e, rs, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// The model is told which engine it searches with.
|
|
func TestSearchPromptNamesEngine(t *testing.T) {
|
|
if p := searchPrompt(searchConf{Engine: "tavily"}); !strings.Contains(p, "Tavily") || strings.Contains(p, "DuckDuckGo") {
|
|
t.Fatalf("tavily: %s", p)
|
|
}
|
|
if p := searchPrompt(searchConf{}); !strings.Contains(p, "DuckDuckGo") {
|
|
t.Fatalf("default: %s", p)
|
|
}
|
|
}
|