110 lines
3.3 KiB
Go
110 lines
3.3 KiB
Go
package io
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
|
)
|
|
|
|
// webSearchToolForTest builds a tools set with opts and returns the web-search
|
|
// tool.
|
|
func webSearchToolForTest(t *testing.T, opts ...Option) *webSearchTool {
|
|
t.Helper()
|
|
return findTool(t, buildTools(t, &config.Config{}, opts...), "web-search").(*webSearchTool)
|
|
}
|
|
|
|
func TestWebSearch_NoEndpoint(t *testing.T) {
|
|
t.Parallel()
|
|
ws := webSearchToolForTest(t) // no endpoint configured
|
|
_, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
|
if err == nil || !strings.Contains(err.Error(), "no search endpoint") {
|
|
t.Fatalf("expected no-endpoint error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestWebSearch_Query(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Query().Get("q") != "cyrene" {
|
|
t.Errorf("query param q = %q", r.URL.Query().Get("q"))
|
|
}
|
|
if r.URL.Query().Get("format") != "json" {
|
|
t.Errorf("format = %q", r.URL.Query().Get("format"))
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"results":[
|
|
{"title":"A","url":"http://a.example","content":"alpha"},
|
|
{"title":"B","url":"http://b.example","snippet":"beta"}
|
|
]}`)
|
|
}))
|
|
defer srv.Close()
|
|
ws := webSearchToolForTest(t,
|
|
WithSearchEndpoint(srv.URL),
|
|
WithAllowedHosts(serverHost(srv)),
|
|
WithHTTPClient(srv.Client()),
|
|
)
|
|
|
|
got, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene", "count": 2})
|
|
if err != nil {
|
|
t.Fatalf("Execute: %v", err)
|
|
}
|
|
m := got.(map[string]any)
|
|
if m["count"].(int) != 2 {
|
|
t.Fatalf("count = %v", m["count"])
|
|
}
|
|
results := m["results"].([]searchResult)
|
|
if results[0].Title != "A" || results[0].Content != "alpha" {
|
|
t.Fatalf("result[0] = %+v", results[0])
|
|
}
|
|
if results[1].Content != "beta" {
|
|
t.Fatalf("result[1] snippet fallback = %+v", results[1])
|
|
}
|
|
}
|
|
|
|
func TestWebSearch_SizeLimit(t *testing.T) {
|
|
t.Parallel()
|
|
big := strings.Repeat("a", 64)
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, _ = io.WriteString(w, big)
|
|
}))
|
|
defer srv.Close()
|
|
ws := webSearchToolForTest(t,
|
|
WithSearchEndpoint(srv.URL),
|
|
WithAllowedHosts(serverHost(srv)),
|
|
WithHTTPClient(srv.Client()),
|
|
WithHTTPResponseLimit(16),
|
|
)
|
|
_, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
|
if err == nil || !strings.Contains(err.Error(), "exceeds the 16-byte limit") {
|
|
t.Fatalf("expected size-limit error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestWebSearch_CountCapped(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"results":[{"title":"A","url":"http://a","content":"x"}]}`)
|
|
}))
|
|
defer srv.Close()
|
|
ws := webSearchToolForTest(t,
|
|
WithSearchEndpoint(srv.URL),
|
|
WithAllowedHosts(serverHost(srv)),
|
|
WithHTTPClient(srv.Client()),
|
|
WithMaxResults(0), // ignored; default 10
|
|
WithMaxResults(1),
|
|
)
|
|
got, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
|
if err != nil {
|
|
t.Fatalf("Execute: %v", err)
|
|
}
|
|
if m := got.(map[string]any); m["count"].(int) != 1 {
|
|
t.Fatalf("count = %v", m["count"])
|
|
}
|
|
}
|