feat: add unit tests for routing strategies and implement dynamic selector updates

Added comprehensive tests for `FillFirstSelector` and `RoundRobinSelector` to ensure proper behavior, including deterministic, cyclical, and concurrent scenarios. Introduced dynamic routing strategy updates in `service.go`, normalizing strategies and seamlessly switching between `fill-first` and `round-robin`. Updated `Manager` to support selector changes via the new `SetSelector` method.
This commit is contained in:
Luis Pater
2025-12-22 22:52:23 +08:00
parent e19ddb53e7
commit b84ccc6e7a
4 changed files with 159 additions and 3 deletions

View File

@@ -135,6 +135,18 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager {
}
}
func (m *Manager) SetSelector(selector Selector) {
if m == nil {
return
}
if selector == nil {
selector = &RoundRobinSelector{}
}
m.mu.Lock()
m.selector = selector
m.mu.Unlock()
}
// SetStore swaps the underlying persistence store.
func (m *Manager) SetStore(store Store) {
m.mu.Lock()

View File

@@ -149,9 +149,6 @@ func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([]
func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
_ = ctx
_ = opts
if s.cursors == nil {
s.cursors = make(map[string]int)
}
now := time.Now()
available, err := getAvailableAuths(auths, provider, model, now)
if err != nil {
@@ -159,6 +156,9 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o
}
key := provider + ":" + model
s.mu.Lock()
if s.cursors == nil {
s.cursors = make(map[string]int)
}
index := s.cursors[key]
if index >= 2_147_483_640 {

View File

@@ -0,0 +1,113 @@
package auth
import (
"context"
"errors"
"sync"
"testing"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
)
func TestFillFirstSelectorPick_Deterministic(t *testing.T) {
t.Parallel()
selector := &FillFirstSelector{}
auths := []*Auth{
{ID: "b"},
{ID: "a"},
{ID: "c"},
}
got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
if err != nil {
t.Fatalf("Pick() error = %v", err)
}
if got == nil {
t.Fatalf("Pick() auth = nil")
}
if got.ID != "a" {
t.Fatalf("Pick() auth.ID = %q, want %q", got.ID, "a")
}
}
func TestRoundRobinSelectorPick_CyclesDeterministic(t *testing.T) {
t.Parallel()
selector := &RoundRobinSelector{}
auths := []*Auth{
{ID: "b"},
{ID: "a"},
{ID: "c"},
}
want := []string{"a", "b", "c", "a", "b"}
for i, id := range want {
got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
if err != nil {
t.Fatalf("Pick() #%d error = %v", i, err)
}
if got == nil {
t.Fatalf("Pick() #%d auth = nil", i)
}
if got.ID != id {
t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, id)
}
}
}
func TestRoundRobinSelectorPick_Concurrent(t *testing.T) {
selector := &RoundRobinSelector{}
auths := []*Auth{
{ID: "b"},
{ID: "a"},
{ID: "c"},
}
start := make(chan struct{})
var wg sync.WaitGroup
errCh := make(chan error, 1)
goroutines := 32
iterations := 100
for i := 0; i < goroutines; i++ {
wg.Add(1)
go func() {
defer wg.Done()
<-start
for j := 0; j < iterations; j++ {
got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
if err != nil {
select {
case errCh <- err:
default:
}
return
}
if got == nil {
select {
case errCh <- errors.New("Pick() returned nil auth"):
default:
}
return
}
if got.ID == "" {
select {
case errCh <- errors.New("Pick() returned auth with empty ID"):
default:
}
return
}
}
}()
}
close(start)
wg.Wait()
select {
case err := <-errCh:
t.Fatalf("concurrent Pick() error = %v", err)
default:
}
}

View File

@@ -506,6 +506,13 @@ func (s *Service) Run(ctx context.Context) error {
var watcherWrapper *WatcherWrapper
reloadCallback := func(newCfg *config.Config) {
previousStrategy := ""
s.cfgMu.RLock()
if s.cfg != nil {
previousStrategy = strings.ToLower(strings.TrimSpace(s.cfg.Routing.Strategy))
}
s.cfgMu.RUnlock()
if newCfg == nil {
s.cfgMu.RLock()
newCfg = s.cfg
@@ -514,6 +521,30 @@ func (s *Service) Run(ctx context.Context) error {
if newCfg == nil {
return
}
nextStrategy := strings.ToLower(strings.TrimSpace(newCfg.Routing.Strategy))
normalizeStrategy := func(strategy string) string {
switch strategy {
case "fill-first", "fillfirst", "ff":
return "fill-first"
default:
return "round-robin"
}
}
previousStrategy = normalizeStrategy(previousStrategy)
nextStrategy = normalizeStrategy(nextStrategy)
if s.coreManager != nil && previousStrategy != nextStrategy {
var selector coreauth.Selector
switch nextStrategy {
case "fill-first":
selector = &coreauth.FillFirstSelector{}
default:
selector = &coreauth.RoundRobinSelector{}
}
s.coreManager.SetSelector(selector)
log.Infof("routing strategy updated to %s", nextStrategy)
}
s.applyRetryConfig(newCfg)
if s.server != nil {
s.server.UpdateClients(newCfg)