Agent-Native/runtime/tools/search_test.go

package tools

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

type fakeProvider struct {
	resp SearchResponse
	err  error
	got  SearchRequest
}

func (f *fakeProvider) Search(_ context.Context, req SearchRequest) (SearchResponse, error) {
	f.got = req
	return f.resp, f.err
}

func newSearchRegistry(p SearchProvider) *Registry {
	r := NewRegistry(tTempDir())
	RegisterSearch(r, p)
	return r
}

func tTempDir() string { return "/tmp/agent-native-search-test" }

func TestSearchToolSchemaAndAlias(t *testing.T) {
	r := newSearchRegistry(&fakeProvider{})
	names := map[string]bool{}
	for _, d := range r.Decls() {
		names[d.Name] = true
	}
	if !names["tavily_search"] || !names["web_search"] {
		t.Fatalf("missing tool decls: %v", names)
	}
}

func TestSearchToolFormatsResults(t *testing.T) {
	fake := &fakeProvider{resp: SearchResponse{
		Answer: "实时总结 https://example.com/a",
		Results: []SearchResult{
			{Title: "标题一", Content: strings.Repeat("摘", 400), Score: 0.91234},
			{Title: "", Content: "  ", Score: 0},
		},
	}}
	r := newSearchRegistry(fake)
	res := r.Execute(context.Background(), "tavily_search", mustJSON(t, map[string]any{"query": "超声 法规"}))
	if res.IsError {
		t.Fatalf("unexpected error: %s", res.Content)
	}
	for _, want := range []string{
		`## 🌐 Tavily 实时网络搜索结果: "超声 法规"`,
		"> **智能总结**: 实时总结 #", // URL 被脱敏为 #
		"### 检索到的参考条目 (2 条):",
		"无标题网页",
		"（无摘要）",
		"**相关度评分**: 0.912",
	} {
		if !strings.Contains(res.Content, want) {
			t.Fatalf("output missing %q:\n%s", want, res.Content)
		}
	}
	// 摘要截断：301 个连续字符不得出现
	if strings.Contains(res.Content, strings.Repeat("摘", 301)) {
		t.Fatal("snippet not truncated at 300")
	}
	if fake.got.Query != "超声 法规" || fake.got.MaxResults != 0 {
		t.Fatalf("request shape: %+v", fake.got)
	}
}

func TestSearchToolSanitizesSensitiveTerms(t *testing.T) {
	fake := &fakeProvider{resp: SearchResponse{Answer: "涉及避孕内容"}}
	r := newSearchRegistry(fake)
	res := r.Execute(context.Background(), "web_search", mustJSON(t, map[string]any{"query": "x"}))
	if !strings.Contains(res.Content, "[医疗]") || strings.Contains(res.Content, "避孕") {
		t.Fatalf("sanitize failed: %s", res.Content)
	}
}

func TestSearchToolEmptyQueryRejected(t *testing.T) {
	r := newSearchRegistry(&fakeProvider{})
	// registry schema 层先拒（required 校验空串）；直调工具层同语义
	res := r.Execute(context.Background(), "tavily_search", mustJSON(t, map[string]any{"query": " "}))
	if !res.IsError || !strings.Contains(res.Content, "query") {
		t.Fatalf("empty query must error: %+v", res)
	}
	direct := (&tavilySearchTool{provider: &fakeProvider{}}).Execute
	if out, err := direct(context.Background(), mustJSON(t, map[string]any{"query": ""})); err != nil || !out.IsError || !strings.Contains(out.Content, "query cannot be empty") {
		t.Fatalf("tool-level empty query: %+v %v", out, err)
	}
}

func TestSearchToolProviderFailureDegrades(t *testing.T) {
	r := newSearchRegistry(&fakeProvider{err: context.DeadlineExceeded})
	res := r.Execute(context.Background(), "tavily_search", mustJSON(t, map[string]any{"query": "q"}))
	if !res.IsError {
		t.Fatal("provider failure must be marked error")
	}
	if !strings.Contains(res.Content, "[Tavily Search Failed]") ||
		!strings.Contains(res.Content, "fallback") {
		t.Fatalf("degradation text missing: %s", res.Content)
	}
}

func TestTavilyProviderKeyRotation(t *testing.T) {
	var calls []string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]any
		_ = json.NewDecoder(r.Body).Decode(&body)
		calls = append(calls, body["api_key"].(string))
		switch body["api_key"] {
		case "key1":
			w.WriteHeader(http.StatusTooManyRequests) // 429 → 轮换
		case "key2":
			w.WriteHeader(http.StatusUnauthorized) // 401 → 轮换
		default:
			_, _ = w.Write([]byte(`{"answer":"ok","results":[{"title":"t","content":"c","score":0.5}]}`))
		}
	}))
	defer srv.Close()

	p := NewTavilyProvider(srv.URL, []string{"key1", "key2", "key3"})
	resp, err := p.Search(context.Background(), SearchRequest{Query: "q"})
	if err != nil {
		t.Fatalf("rotation search: %v", err)
	}
	if resp.Answer != "ok" || len(resp.Results) != 1 {
		t.Fatalf("response: %+v", resp)
	}
	if len(calls) != 3 || calls[0] != "key1" || calls[2] != "key3" {
		t.Fatalf("rotation order: %v", calls)
	}
}

func TestTavilyProviderAllKeysFail(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusForbidden)
	}))
	defer srv.Close()
	p := NewTavilyProvider(srv.URL, []string{"k1", "k2"})
	if _, err := p.Search(context.Background(), SearchRequest{Query: "q"}); err == nil {
		t.Fatal("all keys failing must error")
	}
}

func TestTavilyProviderNoKeys(t *testing.T) {
	p := NewTavilyProvider("", nil)
	if _, err := p.Search(context.Background(), SearchRequest{Query: "q"}); err == nil ||
		!strings.Contains(err.Error(), "no API key") {
		t.Fatalf("no-key must fail closed: %v", err)
	}
}

func mustJSON(t *testing.T, v any) json.RawMessage {
	t.Helper()
	b, err := json.Marshal(v)
	if err != nil {
		t.Fatal(err)
	}
	return b
}