91 lines
3.2 KiB
Go
91 lines
3.2 KiB
Go
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])
|
|
}
|
|
}
|