初始化提交

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

@@ -0,0 +1,123 @@
+package balancer
+
+import (
+	"testing"
+	"time"
+
+	"github.com/chenxin/hunjing/mcp/web-search/internal/models"
+)
+
+func TestKeyPool_Select_Basic(t *testing.T) {
+	kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 5, 30*time.Second)
+	ks := kp.Select()
+	if ks == nil {
+		t.Fatal("expected to select a key")
+	}
+	if ks.Key != "key1" && ks.Key != "key2" {
+		t.Errorf("unexpected key: %s", ks.Key)
+	}
+}
+
+func TestKeyPool_Select_SkipsDead(t *testing.T) {
+	kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 2, 1*time.Hour)
+	ks1 := kp.Select()
+
+	// Fail twice to trigger circuit breaker
+	ks1.RecordFailure()
+	ks1.RecordFailure()
+
+	selected := kp.Select()
+	if selected == nil {
+		t.Fatal("expected fallback to second key")
+	}
+	if selected.Key == ks1.Key {
+		t.Error("should not select the failed key")
+	}
+}
+
+func TestKeyState_CombinedScore(t *testing.T) {
+	ks := newKeyState(models.ProviderTavily, "test-key", 2, 5, 30*time.Second)
+	score := ks.CombinedScore()
+	if score <= 0 {
+		t.Errorf("expected positive score, got %f", score)
+	}
+
+	// After recording a slow latency, score should drop
+	ks.RecordSuccess(2 * time.Second)
+	slowScore := ks.CombinedScore()
+	if slowScore >= score {
+		t.Error("expected slow latency to reduce score")
+	}
+}
+
+func TestCircuitBreaker_OpenAndRecover(t *testing.T) {
+	cb := NewCircuitBreaker(3, 50*time.Millisecond)
+
+	// Should be closed initially
+	if !cb.Allow() {
+		t.Error("expected allow initially")
+	}
+
+	// Fail 3 times to open
+	cb.Failure()
+	cb.Failure()
+	cb.Failure()
+	if cb.Allow() {
+		t.Error("expected circuit open after 3 failures")
+	}
+
+	// Wait for cooldown
+	time.Sleep(60 * time.Millisecond)
+	if !cb.Allow() {
+		t.Error("expected allow after cooldown")
+	}
+
+	// Half-open: success should close
+	cb.Success()
+	if cb.State() != StateClosed {
+		t.Error("expected circuit closed after half-open success")
+	}
+}
+
+func TestCircuitBreaker_HalfOpenFailure(t *testing.T) {
+	cb := NewCircuitBreaker(3, 50*time.Millisecond)
+
+	cb.Failure()
+	cb.Failure()
+	cb.Failure()
+	if cb.Allow() {
+		t.Error("expected open")
+	}
+
+	time.Sleep(60 * time.Millisecond)
+	if !cb.Allow() {
+		t.Error("expected half-open allow")
+	}
+
+	// Half-open: failure should re-open immediately
+	cb.Failure()
+	if cb.State() != StateOpen {
+		t.Error("expected open after half-open failure")
+	}
+}
+
+func TestKeyPool_ActiveCount(t *testing.T) {
+	kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2", "key3"}, 1, 3, 30*time.Second)
+	if kp.TotalCount() != 3 {
+		t.Errorf("expected 3 total, got %d", kp.TotalCount())
+	}
+	if kp.ActiveCount() != 3 {
+		t.Errorf("expected 3 active, got %d", kp.ActiveCount())
+	}
+}
+
+func TestKeyPool_AllKeysDead(t *testing.T) {
+	kp := NewKeyPool(models.ProviderTavily, []string{"key1"}, 1, 1, 1*time.Hour)
+	ks := kp.Select()
+	ks.RecordFailure()
+
+	selected := kp.Select()
+	if selected != nil {
+		t.Error("expected nil when all keys dead")
+	}
+}