thread-master/backend/internal/library/knowledge/brave_collect_test.go

package knowledge

import (
	"context"
	"fmt"
	"sync"
	"testing"

	"haixun-backend/internal/library/websearch"
)

func TestUniqueSourceCount(t *testing.T) {
	count := uniqueSourceCount([]BraveSource{
		{URL: "https://a.example"},
		{URL: "https://a.example"},
		{URL: "https://b.example"},
		{URL: ""},
	})
	if count != 2 {
		t.Fatalf("expected 2 unique urls, got %d", count)
	}
}

func TestSourcesAreSufficientRequiresProcessedQuality(t *testing.T) {
	items := make([]BraveSource, 0, 14)
	for i := 0; i < 14; i++ {
		items = append(items, BraveSource{
			Query: "q1",
			URL:   fmt.Sprintf("https://example.com/%d", i),
			Title: "薄",
		})
	}
	cfg := BraveCollectConfig{MinSourcesBeforeStop: 14, MaxSourcesCap: 22}
	if sourcesAreSufficient(items, []string{"q1", "q2"}, cfg) {
		t.Fatal("thin same-query results should not be considered sufficient")
	}

	for i := range items {
		items[i].Snippet = "這是一段整理後足以支撐判斷的搜尋摘要內容，包含痛點與情境"
		if i >= 7 {
			items[i].Query = "q2"
		}
	}
	if !sourcesAreSufficient(items, []string{"q1", "q2"}, cfg) {
		t.Fatal("diverse useful results should be considered sufficient")
	}
}

func TestCollectWebSourcesParallelStopsAfterSufficientBatch(t *testing.T) {
	client := &fakeSearchClient{results: map[string][]websearch.SearchResult{}}
	for qi := 0; qi < 5; qi++ {
		query := fmt.Sprintf("q%d", qi+1)
		for ri := 0; ri < 7; ri++ {
			client.results[query] = append(client.results[query], websearch.SearchResult{
				URL:     fmt.Sprintf("https://example.com/%s/%d", query, ri),
				Title:   "有用來源",
				Snippet: "這是一段整理後足以支撐判斷的搜尋摘要內容，包含痛點與情境",
			})
		}
	}

	out := CollectWebSources(
		context.Background(),
		client,
		BraveSearchLocale{},
		[]string{"q1", "q2", "q3", "q4", "q5"},
		BraveCollectConfig{ResultsPerQuery: 7, MinSourcesBeforeStop: 14, MaxSourcesCap: 30, Concurrency: 2},
		nil,
		nil,
	)
	if len(out) != 14 {
		t.Fatalf("expected first batch to provide 14 sources, got %d", len(out))
	}
	if got := client.callCount(); got != 2 {
		t.Fatalf("expected collection to stop after first sufficient batch, called %d queries", got)
	}
}

func TestPlanQueriesHybridBudget(t *testing.T) {
	queries := PlanQueries(PlanInput{
		Seed:           "敏感肌",
		TargetAudience: "孕婦",
		Questions:      []string{"q1", "q2", "q3", "q4"},
		Pillars:        []string{"p1", "p2", "p3"},
		L1Labels:       []string{"a", "b", "c", "d"},
		Strategy:       ExpandStrategyHybrid,
	})
	if len(queries) > 4 {
		t.Fatalf("hybrid should cap at 4 queries, got %d: %v", len(queries), queries)
	}
	if len(queries) == 0 {
		t.Fatal("expected some queries")
	}
}

func TestPlanQueriesPrioritizesPatrolKeywords(t *testing.T) {
	queries := PlanQueries(PlanInput{
		Seed:           "敏感肌",
		PatrolKeywords: []string{"化療 沐浴乳", "無香沐浴乳 推薦"},
		Questions:      []string{"很長的受眾提問一", "很長的受眾提問二"},
		Strategy:       ExpandStrategyBrave,
	})
	if len(queries) == 0 {
		t.Fatal("expected queries")
	}
	if queries[0] != "化療 沐浴乳" {
		t.Fatalf("expected patrol keyword first, got %q", queries[0])
	}
}

func TestQueriesExcept(t *testing.T) {
	remaining := QueriesExcept(
		[]string{"a", "b", "c", "d"},
		[]string{"a", "c"},
	)
	if len(remaining) != 2 || remaining[0] != "b" || remaining[1] != "d" {
		t.Fatalf("unexpected remaining: %v", remaining)
	}
}

func TestPlanBootstrapQueriesSkipsResearchMapFields(t *testing.T) {
	full := PlanQueries(PlanInput{
		Seed:      "敏感肌",
		Questions: []string{"q1"},
		Pillars:   []string{"p1"},
		Strategy:  ExpandStrategyBrave,
	})
	bootstrap := PlanBootstrapQueries(PlanInput{
		Seed:      "敏感肌",
		Questions: []string{"q1"},
		Pillars:   []string{"p1"},
		Strategy:  ExpandStrategyBrave,
	})
	if len(bootstrap) == 0 {
		t.Fatal("expected bootstrap queries")
	}
	for _, q := range bootstrap {
		if q == "q1" || q == "p1 請問" {
			t.Fatalf("bootstrap should not include research map query %q", q)
		}
	}
	if len(full) <= len(bootstrap) {
		t.Fatalf("full plan should include more queries than bootstrap: full=%v bootstrap=%v", full, bootstrap)
	}
}

func TestMergeBraveSourcesDedupesURLs(t *testing.T) {
	merged := MergeBraveSources(
		[]BraveSource{{URL: "https://a.example", Query: "q1"}},
		[]BraveSource{{URL: "https://a.example", Query: "q2"}, {URL: "https://b.example", Query: "q3"}},
	)
	if len(merged) != 2 {
		t.Fatalf("expected 2 merged sources, got %d", len(merged))
	}
}

func TestSupplementalQueriesSkippedForHybrid(t *testing.T) {
	queries := PlanQueries(PlanInput{
		Seed:         "敏感肌",
		Supplemental: true,
		Strategy:     ExpandStrategyHybrid,
	})
	if len(queries) != 0 {
		t.Fatalf("hybrid supplemental brave should be empty, got %v", queries)
	}
}

type fakeSearchClient struct {
	mu      sync.Mutex
	calls   []string
	results map[string][]websearch.SearchResult
}

func (f *fakeSearchClient) Enabled() bool { return true }

func (f *fakeSearchClient) Provider() websearch.Provider { return websearch.ProviderBrave }

func (f *fakeSearchClient) Search(_ context.Context, opts websearch.SearchOptions) (websearch.SearchResponse, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.results == nil {
		f.results = map[string][]websearch.SearchResult{}
	}
	f.calls = append(f.calls, opts.Query)
	return websearch.SearchResponse{
		Query:    opts.Query,
		Status:   "success",
		Provider: websearch.ProviderBrave,
		Results:  append([]websearch.SearchResult(nil), f.results[opts.Query]...),
	}, nil
}

func (f *fakeSearchClient) callCount() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return len(f.calls)
}