VooDisss
6962e09dd9
Keep sticky auth affinity limited to matching providers and stop persisting execution-session IDs as long-lived affinity keys so provider switching and normal streaming traffic do not create incorrect pins or stale affinity state.
101 lines
3.7 KiB
Go
101 lines
3.7 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"testing"
|
|
|
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
|
)
|
|
|
|
type affinityTestExecutor struct{ id string }
|
|
|
|
func (e affinityTestExecutor) Identifier() string { return e.id }
|
|
|
|
func (e affinityTestExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
|
|
func (e affinityTestExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
ch := make(chan cliproxyexecutor.StreamChunk)
|
|
close(ch)
|
|
return &cliproxyexecutor.StreamResult{Chunks: ch}, nil
|
|
}
|
|
|
|
func (e affinityTestExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { return auth, nil }
|
|
|
|
func (e affinityTestExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
|
|
func (e affinityTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestManagerPickNextMixedUsesAuthAffinity(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
|
manager.executors["codex"] = affinityTestExecutor{id: "codex"}
|
|
reg := registry.GetGlobalRegistry()
|
|
reg.RegisterClient("codex-a", "codex", []*registry.ModelInfo{{ID: "gpt-5.4"}})
|
|
reg.RegisterClient("codex-b", "codex", []*registry.ModelInfo{{ID: "gpt-5.4"}})
|
|
t.Cleanup(func() {
|
|
reg.UnregisterClient("codex-a")
|
|
reg.UnregisterClient("codex-b")
|
|
})
|
|
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "codex-a", Provider: "codex"}); errRegister != nil {
|
|
t.Fatalf("Register(codex-a) error = %v", errRegister)
|
|
}
|
|
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "codex-b", Provider: "codex"}); errRegister != nil {
|
|
t.Fatalf("Register(codex-b) error = %v", errRegister)
|
|
}
|
|
|
|
manager.SetAuthAffinity("codex", "idem-1", "codex-b")
|
|
opts := cliproxyexecutor.Options{Metadata: map[string]any{"auth_affinity_key": "idem-1"}}
|
|
|
|
got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"codex"}, "gpt-5.4", opts, map[string]struct{}{})
|
|
if errPick != nil {
|
|
t.Fatalf("pickNextMixed() error = %v", errPick)
|
|
}
|
|
if provider != "codex" {
|
|
t.Fatalf("provider = %q, want %q", provider, "codex")
|
|
}
|
|
if got == nil || got.ID != "codex-b" {
|
|
t.Fatalf("auth.ID = %v, want codex-b", got)
|
|
}
|
|
if pinned := pinnedAuthIDFromMetadata(opts.Metadata); pinned != "codex-b" {
|
|
t.Fatalf("pinned auth metadata = %q, want %q", pinned, "codex-b")
|
|
}
|
|
}
|
|
|
|
func TestManagerAuthAffinityRoundTrip(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetAuthAffinity("codex", "idem-2", "auth-1")
|
|
if got := manager.AuthAffinity("codex", "idem-2"); got != "auth-1" {
|
|
t.Fatalf("AuthAffinity = %q, want %q", got, "auth-1")
|
|
}
|
|
manager.ClearAuthAffinity("codex", "idem-2")
|
|
if got := manager.AuthAffinity("codex", "idem-2"); got != "" {
|
|
t.Fatalf("AuthAffinity after clear = %q, want empty", got)
|
|
}
|
|
}
|
|
|
|
func TestManagerAuthAffinityScopedByProvider(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetAuthAffinity("codex", "shared-key", "codex-auth")
|
|
manager.SetAuthAffinity("gemini", "shared-key", "gemini-auth")
|
|
|
|
if got := manager.AuthAffinity("codex", "shared-key"); got != "codex-auth" {
|
|
t.Fatalf("codex affinity = %q, want %q", got, "codex-auth")
|
|
}
|
|
if got := manager.AuthAffinity("gemini", "shared-key"); got != "gemini-auth" {
|
|
t.Fatalf("gemini affinity = %q, want %q", got, "gemini-auth")
|
|
}
|
|
}
|