253 lines
7 KiB
Go
253 lines
7 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestManager_RegisterCanonicalizesThinkingSuffixModelStates(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
now := time.Now()
|
|
laterRetry := now.Add(2 * time.Hour)
|
|
|
|
registered, errRegister := manager.Register(context.Background(), &Auth{
|
|
ID: "auth-thinking-states",
|
|
Provider: "gemini",
|
|
ModelStates: map[string]*ModelState{
|
|
"gemini-3.1-pro-preview(high)": {
|
|
Status: StatusError,
|
|
Unavailable: true,
|
|
NextRetryAfter: now.Add(time.Hour),
|
|
Quota: QuotaState{
|
|
Exceeded: true,
|
|
NextRecoverAt: now.Add(time.Hour),
|
|
BackoffLevel: 1,
|
|
},
|
|
UpdatedAt: now,
|
|
},
|
|
"gemini-3.1-pro-preview(low)": {
|
|
Status: StatusError,
|
|
Unavailable: true,
|
|
NextRetryAfter: laterRetry,
|
|
Quota: QuotaState{
|
|
Exceeded: true,
|
|
NextRecoverAt: laterRetry,
|
|
BackoffLevel: 2,
|
|
},
|
|
UpdatedAt: now.Add(time.Minute),
|
|
},
|
|
},
|
|
})
|
|
if errRegister != nil {
|
|
t.Fatalf("Register() error = %v", errRegister)
|
|
}
|
|
if len(registered.ModelStates) != 1 {
|
|
t.Fatalf("len(ModelStates) = %d, want 1: %+v", len(registered.ModelStates), registered.ModelStates)
|
|
}
|
|
state := registered.ModelStates["gemini-3.1-pro-preview"]
|
|
if state == nil || !state.Unavailable || !state.NextRetryAfter.Equal(laterRetry) {
|
|
t.Fatalf("canonical model state = %+v, want unavailable until %v", state, laterRetry)
|
|
}
|
|
if state.Quota.BackoffLevel != 2 || !state.Quota.NextRecoverAt.Equal(laterRetry) {
|
|
t.Fatalf("canonical model quota = %+v, want latest cooldown", state.Quota)
|
|
}
|
|
}
|
|
|
|
func TestManager_Update_PreservesModelStates(t *testing.T) {
|
|
m := NewManager(nil, nil, nil)
|
|
|
|
model := "test-model"
|
|
backoffLevel := 7
|
|
|
|
if _, errRegister := m.Register(context.Background(), &Auth{
|
|
ID: "auth-1",
|
|
Provider: "claude",
|
|
Metadata: map[string]any{"k": "v"},
|
|
ModelStates: map[string]*ModelState{
|
|
model: {
|
|
Quota: QuotaState{BackoffLevel: backoffLevel},
|
|
},
|
|
},
|
|
}); errRegister != nil {
|
|
t.Fatalf("register auth: %v", errRegister)
|
|
}
|
|
|
|
if _, errUpdate := m.Update(context.Background(), &Auth{
|
|
ID: "auth-1",
|
|
Provider: "claude",
|
|
Metadata: map[string]any{"k": "v2"},
|
|
}); errUpdate != nil {
|
|
t.Fatalf("update auth: %v", errUpdate)
|
|
}
|
|
|
|
updated, ok := m.GetByID("auth-1")
|
|
if !ok || updated == nil {
|
|
t.Fatalf("expected auth to be present")
|
|
}
|
|
if len(updated.ModelStates) == 0 {
|
|
t.Fatalf("expected ModelStates to be preserved")
|
|
}
|
|
state := updated.ModelStates[model]
|
|
if state == nil {
|
|
t.Fatalf("expected model state to be present")
|
|
}
|
|
if state.Quota.BackoffLevel != backoffLevel {
|
|
t.Fatalf("expected BackoffLevel to be %d, got %d", backoffLevel, state.Quota.BackoffLevel)
|
|
}
|
|
}
|
|
|
|
func TestManager_Update_DisabledExistingDoesNotInheritModelStates(t *testing.T) {
|
|
m := NewManager(nil, nil, nil)
|
|
|
|
// Register a disabled auth with existing ModelStates.
|
|
if _, err := m.Register(context.Background(), &Auth{
|
|
ID: "auth-disabled",
|
|
Provider: "claude",
|
|
Disabled: true,
|
|
Status: StatusDisabled,
|
|
ModelStates: map[string]*ModelState{
|
|
"stale-model": {
|
|
Quota: QuotaState{BackoffLevel: 5},
|
|
},
|
|
},
|
|
}); err != nil {
|
|
t.Fatalf("register auth: %v", err)
|
|
}
|
|
|
|
// Update with empty ModelStates — should NOT inherit stale states.
|
|
if _, err := m.Update(context.Background(), &Auth{
|
|
ID: "auth-disabled",
|
|
Provider: "claude",
|
|
Disabled: true,
|
|
Status: StatusDisabled,
|
|
}); err != nil {
|
|
t.Fatalf("update auth: %v", err)
|
|
}
|
|
|
|
updated, ok := m.GetByID("auth-disabled")
|
|
if !ok || updated == nil {
|
|
t.Fatalf("expected auth to be present")
|
|
}
|
|
if len(updated.ModelStates) != 0 {
|
|
t.Fatalf("expected disabled auth NOT to inherit ModelStates, got %d entries", len(updated.ModelStates))
|
|
}
|
|
}
|
|
|
|
func TestManager_Update_ActiveToDisabledDoesNotInheritModelStates(t *testing.T) {
|
|
m := NewManager(nil, nil, nil)
|
|
|
|
// Register an active auth with ModelStates (simulates existing live auth).
|
|
if _, err := m.Register(context.Background(), &Auth{
|
|
ID: "auth-a2d",
|
|
Provider: "claude",
|
|
Status: StatusActive,
|
|
ModelStates: map[string]*ModelState{
|
|
"stale-model": {
|
|
Quota: QuotaState{BackoffLevel: 9},
|
|
},
|
|
},
|
|
}); err != nil {
|
|
t.Fatalf("register auth: %v", err)
|
|
}
|
|
|
|
// File watcher deletes config → synthesizes Disabled=true auth → Update.
|
|
// Even though existing is active, incoming auth is disabled → skip inheritance.
|
|
if _, err := m.Update(context.Background(), &Auth{
|
|
ID: "auth-a2d",
|
|
Provider: "claude",
|
|
Disabled: true,
|
|
Status: StatusDisabled,
|
|
}); err != nil {
|
|
t.Fatalf("update auth: %v", err)
|
|
}
|
|
|
|
updated, ok := m.GetByID("auth-a2d")
|
|
if !ok || updated == nil {
|
|
t.Fatalf("expected auth to be present")
|
|
}
|
|
if len(updated.ModelStates) != 0 {
|
|
t.Fatalf("expected active→disabled transition NOT to inherit ModelStates, got %d entries", len(updated.ModelStates))
|
|
}
|
|
}
|
|
|
|
func TestManager_Update_DisabledToActiveDoesNotInheritStaleModelStates(t *testing.T) {
|
|
m := NewManager(nil, nil, nil)
|
|
|
|
// Register a disabled auth with stale ModelStates.
|
|
if _, err := m.Register(context.Background(), &Auth{
|
|
ID: "auth-d2a",
|
|
Provider: "claude",
|
|
Disabled: true,
|
|
Status: StatusDisabled,
|
|
ModelStates: map[string]*ModelState{
|
|
"stale-model": {
|
|
Quota: QuotaState{BackoffLevel: 4},
|
|
},
|
|
},
|
|
}); err != nil {
|
|
t.Fatalf("register auth: %v", err)
|
|
}
|
|
|
|
// Re-enable: incoming auth is active, existing is disabled → skip inheritance.
|
|
if _, err := m.Update(context.Background(), &Auth{
|
|
ID: "auth-d2a",
|
|
Provider: "claude",
|
|
Status: StatusActive,
|
|
}); err != nil {
|
|
t.Fatalf("update auth: %v", err)
|
|
}
|
|
|
|
updated, ok := m.GetByID("auth-d2a")
|
|
if !ok || updated == nil {
|
|
t.Fatalf("expected auth to be present")
|
|
}
|
|
if len(updated.ModelStates) != 0 {
|
|
t.Fatalf("expected disabled→active transition NOT to inherit stale ModelStates, got %d entries", len(updated.ModelStates))
|
|
}
|
|
}
|
|
|
|
func TestManager_Update_ActiveInheritsModelStates(t *testing.T) {
|
|
m := NewManager(nil, nil, nil)
|
|
|
|
model := "active-model"
|
|
backoffLevel := 3
|
|
|
|
// Register an active auth with ModelStates.
|
|
if _, err := m.Register(context.Background(), &Auth{
|
|
ID: "auth-active",
|
|
Provider: "claude",
|
|
Status: StatusActive,
|
|
ModelStates: map[string]*ModelState{
|
|
model: {
|
|
Quota: QuotaState{BackoffLevel: backoffLevel},
|
|
},
|
|
},
|
|
}); err != nil {
|
|
t.Fatalf("register auth: %v", err)
|
|
}
|
|
|
|
// Update with empty ModelStates — both sides active → SHOULD inherit.
|
|
if _, err := m.Update(context.Background(), &Auth{
|
|
ID: "auth-active",
|
|
Provider: "claude",
|
|
Status: StatusActive,
|
|
}); err != nil {
|
|
t.Fatalf("update auth: %v", err)
|
|
}
|
|
|
|
updated, ok := m.GetByID("auth-active")
|
|
if !ok || updated == nil {
|
|
t.Fatalf("expected auth to be present")
|
|
}
|
|
if len(updated.ModelStates) == 0 {
|
|
t.Fatalf("expected active auth to inherit ModelStates")
|
|
}
|
|
state := updated.ModelStates[model]
|
|
if state == nil {
|
|
t.Fatalf("expected model state to be present")
|
|
}
|
|
if state.Quota.BackoffLevel != backoffLevel {
|
|
t.Fatalf("expected BackoffLevel to be %d, got %d", backoffLevel, state.Quota.BackoffLevel)
|
|
}
|
|
}
|