初始化提交

mcp/web-search/internal/balancer/balancer_test.go

@@ -0,0 +1,90 @@
+package balancer
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/chenxin/hunjing/mcp/web-search/internal/config"
+	"github.com/chenxin/hunjing/mcp/web-search/internal/models"
+)
+
+type mockSearcher struct {
+	name   models.ProviderKind
+	failOn int // 第几次调用失败（0 = 永不）
+	calls  int
+}
+
+func (m *mockSearcher) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
+	m.calls++
+	if m.failOn > 0 && m.calls >= m.failOn {
+		return nil, errors.New("mock failure")
+	}
+	return &models.SearchResult{Provider: m.name, Query: req.Query, Results: []models.ResultItem{{Title: "mock"}}}, nil
+}
+func (m *mockSearcher) HealthCheck(ctx context.Context) error { return nil }
+func (m *mockSearcher) Name() models.ProviderKind             { return m.name }
+
+func TestLoadBalancer_Search_Basic(t *testing.T) {
+	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
+	lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})
+
+	result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if result.Provider != models.ProviderTavily {
+		t.Errorf("expected tavily, got %s", result.Provider)
+	}
+}
+
+func TestLoadBalancer_Search_FallbackProvider(t *testing.T) {
+	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
+	lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})
+	lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})
+
+	result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if string(result.Provider) != "brave" {
+		t.Errorf("expected fallback to brave, got %s", result.Provider)
+	}
+}
+
+func TestLoadBalancer_Search_AllExhausted(t *testing.T) {
+	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
+	lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})
+
+	_, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
+	if err == nil {
+		t.Fatal("expected error when all keys exhausted")
+	}
+}
+
+func TestLoadBalancer_Search_PreferredProviders(t *testing.T) {
+	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
+	lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})
+	lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})
+
+	result, err := lb.Search(context.Background(), &models.SearchRequest{
+		Query:     "test",
+		Providers: []string{"brave"},
+	})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if result.Provider != models.ProviderBrave {
+		t.Errorf("expected bing (preferred), got %s", result.Provider)
+	}
+}
+
+func TestLoadBalancer_ActiveKeyCount(t *testing.T) {
+	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
+	lb.AddProvider(models.ProviderTavily, []string{"key1", "key2"}, 1, &mockSearcher{name: models.ProviderTavily})
+
+	counts := lb.ActiveKeyCount()
+	if counts[models.ProviderTavily] != 2 {
+		t.Errorf("expected 2 active keys, got %d", counts[models.ProviderTavily])
+	}
+}