Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
1802 changed files with 503346 additions and 2 deletions
250
backend/sdk/cliproxy/auth/home_selection_test.go
Normal file
250
backend/sdk/cliproxy/auth/home_selection_test.go
Normal file
|
|
@ -0,0 +1,250 @@
|
|||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
)
|
||||
|
||||
func TestHomeDispatchSelectionOwnsScopeOutsideAuth(t *testing.T) {
|
||||
registry := executionregistry.New()
|
||||
pending, errBegin := registry.BeginDispatch()
|
||||
if errBegin != nil {
|
||||
t.Fatal(errBegin)
|
||||
}
|
||||
scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{RequestID: "req-1", CredentialID: "cred-1", Model: "gpt", Kind: "http", StartedAt: time.Now()})
|
||||
if errInstall != nil {
|
||||
t.Fatal(errInstall)
|
||||
}
|
||||
selection, errSelection := newHomeDispatchSelection(&Auth{ID: "cred-1", Provider: "codex"}, nil, "codex", scope)
|
||||
if errSelection != nil {
|
||||
t.Fatal(errSelection)
|
||||
}
|
||||
clone := selection.CloneAuth()
|
||||
if clone == nil || clone.ID != "cred-1" || clone.Runtime != nil {
|
||||
t.Fatalf("clone = %#v", clone)
|
||||
}
|
||||
closed := atomic.Int32{}
|
||||
if errBind := selection.Bind(func() error { closed.Add(1); return nil }); errBind != nil {
|
||||
t.Fatal(errBind)
|
||||
}
|
||||
selection.End("completed")
|
||||
selection.End("duplicate")
|
||||
if closed.Load() != 1 {
|
||||
t.Fatalf("close calls = %d", closed.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeDispatchSelectionReplaceAuthPreservesRoutingAttributes(t *testing.T) {
|
||||
selection := &HomeDispatchSelection{Auth: &Auth{
|
||||
ID: "cred-1",
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{
|
||||
homeUpstreamModelAttributeKey: "gpt-5-upstream",
|
||||
homeForceMappingAttributeKey: "true",
|
||||
homeOriginalAliasAttributeKey: "team/gpt-5",
|
||||
},
|
||||
Metadata: map[string]any{"access_token": "old"},
|
||||
}}
|
||||
|
||||
selection.ReplaceAuth(&Auth{
|
||||
ID: "cred-1",
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{AttributeAuthKind: AuthKindOAuth},
|
||||
Metadata: map[string]any{"access_token": "fresh"},
|
||||
})
|
||||
|
||||
updated := selection.CloneAuth()
|
||||
if updated == nil || updated.Metadata["access_token"] != "fresh" {
|
||||
t.Fatalf("updated auth = %#v", updated)
|
||||
}
|
||||
if updated.Attributes[homeUpstreamModelAttributeKey] != "gpt-5-upstream" || updated.Attributes[homeForceMappingAttributeKey] != "true" || updated.Attributes[homeOriginalAliasAttributeKey] != "team/gpt-5" {
|
||||
t.Fatalf("routing attributes were not preserved: %#v", updated.Attributes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeDispatchSelectionReplaceAuthConcurrentClone(t *testing.T) {
|
||||
selection := &HomeDispatchSelection{Auth: &Auth{ID: "cred-1", Metadata: map[string]any{"access_token": "old"}}}
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
for i := 0; i < 1000; i++ {
|
||||
selection.ReplaceAuth(&Auth{ID: "cred-1", Metadata: map[string]any{"access_token": "fresh"}})
|
||||
}
|
||||
}()
|
||||
for i := 0; i < 1000; i++ {
|
||||
if auth := selection.CloneAuth(); auth == nil || auth.ID != "cred-1" {
|
||||
t.Fatalf("CloneAuth() = %#v", auth)
|
||||
}
|
||||
}
|
||||
<-done
|
||||
}
|
||||
|
||||
func TestReplaceHomeSelectionAuthUpdatesRetainedRuntimeAuth(t *testing.T) {
|
||||
selection := &HomeDispatchSelection{Auth: &Auth{ID: "cred-1", Provider: "codex", Metadata: map[string]any{"access_token": "old"}}}
|
||||
manager := &Manager{
|
||||
homeRuntimeAuths: map[string]map[string]*Auth{
|
||||
"session-1": {"cred-1": selection.Auth.Clone()},
|
||||
},
|
||||
homeRuntimeAuthOwners: map[string]map[string]*HomeDispatchSelection{
|
||||
"session-1": {"cred-1": selection},
|
||||
},
|
||||
}
|
||||
|
||||
manager.replaceHomeSelectionAuth(selection, &Auth{ID: "cred-1", Provider: "codex", Metadata: map[string]any{"access_token": "fresh"}})
|
||||
|
||||
retained := manager.homeRuntimeAuths["session-1"]["cred-1"]
|
||||
if retained == nil || retained.Metadata["access_token"] != "fresh" {
|
||||
t.Fatalf("retained runtime auth = %#v, want fresh token", retained)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeDispatchSelectionDrainsResourcesAddedDuringEnd(t *testing.T) {
|
||||
registry := executionregistry.New()
|
||||
pending, errBegin := registry.BeginDispatch()
|
||||
if errBegin != nil {
|
||||
t.Fatal(errBegin)
|
||||
}
|
||||
scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{})
|
||||
if errInstall != nil {
|
||||
t.Fatal(errInstall)
|
||||
}
|
||||
selection, errSelection := newHomeDispatchSelection(&Auth{ID: "cred-1"}, nil, "test", scope)
|
||||
if errSelection != nil {
|
||||
t.Fatal(errSelection)
|
||||
}
|
||||
|
||||
started := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
if errBind := selection.Bind(func() error {
|
||||
close(started)
|
||||
<-release
|
||||
return nil
|
||||
}); errBind != nil {
|
||||
t.Fatal(errBind)
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
selection.End("draining")
|
||||
close(done)
|
||||
}()
|
||||
<-started
|
||||
|
||||
closedLate := atomic.Int32{}
|
||||
errLate := selection.Bind(func() error {
|
||||
closedLate.Add(1)
|
||||
return errors.New("late close")
|
||||
})
|
||||
if !errors.Is(errLate, executionregistry.ErrRegistryNotAccepting) {
|
||||
t.Fatalf("late Bind() error = %v, want ErrRegistryNotAccepting", errLate)
|
||||
}
|
||||
if closedLate.Load() != 1 {
|
||||
t.Fatalf("late close calls = %d, want 1", closedLate.Load())
|
||||
}
|
||||
|
||||
close(release)
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("End did not complete")
|
||||
}
|
||||
|
||||
drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancelDrain()
|
||||
if errDrain := registry.Drain(drainCtx); errDrain != nil {
|
||||
t.Fatalf("Drain() error = %v", errDrain)
|
||||
}
|
||||
}
|
||||
|
||||
type gatedHomeDispatcher struct {
|
||||
loaded chan struct{}
|
||||
release chan struct{}
|
||||
rpop atomic.Int32
|
||||
}
|
||||
|
||||
func (d *gatedHomeDispatcher) HeartbeatOK() bool {
|
||||
select {
|
||||
case <-d.loaded:
|
||||
default:
|
||||
close(d.loaded)
|
||||
}
|
||||
<-d.release
|
||||
return true
|
||||
}
|
||||
|
||||
func (d *gatedHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {
|
||||
d.rpop.Add(1)
|
||||
return nil, errors.New("old Home dispatcher was used")
|
||||
}
|
||||
|
||||
func (*gatedHomeDispatcher) AbortAmbiguousDispatch() {}
|
||||
|
||||
func TestManagerHomeDispatchBundleCompareAndClearDoesNotRemoveReplacement(t *testing.T) {
|
||||
manager := NewManager(nil, nil, nil)
|
||||
first := manager.PublishHomeDispatch(&gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}, executionregistry.New(), 1)
|
||||
second := manager.PublishHomeDispatch(&gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}, executionregistry.New(), 2)
|
||||
|
||||
if manager.ClearHomeDispatchBundle(first) {
|
||||
t.Fatal("ClearHomeDispatchBundle() cleared a replacement bundle")
|
||||
}
|
||||
if got := manager.HomeDispatchBundle(); got != second {
|
||||
t.Fatalf("HomeDispatchBundle() = %p, want %p", got, second)
|
||||
}
|
||||
if !manager.ClearHomeDispatchBundle(second) {
|
||||
t.Fatal("ClearHomeDispatchBundle() = false, want true")
|
||||
}
|
||||
if got := manager.HomeDispatchBundle(); got != nil {
|
||||
t.Fatalf("HomeDispatchBundle() = %p, want nil", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPickHomeDispatchSelectionDoesNotMixDetachedBundleWithReplacement(t *testing.T) {
|
||||
manager := NewManager(nil, nil, nil)
|
||||
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
||||
oldDispatcher := &gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}
|
||||
oldRegistry := executionregistry.New()
|
||||
oldBundle := manager.PublishHomeDispatch(oldDispatcher, oldRegistry, 1)
|
||||
|
||||
result := make(chan error, 1)
|
||||
go func() {
|
||||
_, errSelect := manager.pickHomeDispatchSelection(context.Background(), "gpt-5.4", cliproxyexecutor.Options{})
|
||||
result <- errSelect
|
||||
}()
|
||||
select {
|
||||
case <-oldDispatcher.loaded:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("selection did not load the old dispatch bundle")
|
||||
}
|
||||
|
||||
if !manager.ClearHomeDispatchBundle(oldBundle) {
|
||||
t.Fatal("ClearHomeDispatchBundle() = false, want true")
|
||||
}
|
||||
drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancelDrain()
|
||||
if errDrain := oldRegistry.Drain(drainCtx); errDrain != nil {
|
||||
t.Fatalf("old registry Drain() error = %v", errDrain)
|
||||
}
|
||||
manager.PublishHomeDispatch(&gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}, executionregistry.New(), 2)
|
||||
close(oldDispatcher.release)
|
||||
|
||||
select {
|
||||
case errSelect := <-result:
|
||||
var authErr *Error
|
||||
if !errors.As(errSelect, &authErr) || authErr.Code != "home_unavailable" {
|
||||
t.Fatalf("pickHomeDispatchSelection() error = %v, want home_unavailable", errSelect)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("selection did not resume after the old bundle was detached")
|
||||
}
|
||||
if got := oldDispatcher.rpop.Load(); got != 0 {
|
||||
t.Fatalf("old dispatcher RPopAuth() calls = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue