124 lines
2.8 KiB
Go
124 lines
2.8 KiB
Go
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")
|
|
}
|
|
}
|