wehub-resource-sync
305 lines
8.9 KiB
Go
305 lines
8.9 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
type recordingCooldownStateStore struct {
|
|
saveCount atomic.Int32
|
|
mu sync.Mutex
|
|
records []CooldownStateRecord
|
|
load []CooldownStateRecord
|
|
}
|
|
|
|
func (s *recordingCooldownStateStore) Load(context.Context) ([]CooldownStateRecord, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
return cloneCooldownStateRecords(s.load), nil
|
|
}
|
|
|
|
func (s *recordingCooldownStateStore) Save(_ context.Context, records []CooldownStateRecord) error {
|
|
s.saveCount.Add(1)
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.records = cloneCooldownStateRecords(records)
|
|
return nil
|
|
}
|
|
|
|
func cloneCooldownStateRecords(records []CooldownStateRecord) []CooldownStateRecord {
|
|
if len(records) == 0 {
|
|
return nil
|
|
}
|
|
cloned := make([]CooldownStateRecord, len(records))
|
|
for i := range records {
|
|
cloned[i] = records[i]
|
|
cloned[i].LastError = cloneError(records[i].LastError)
|
|
}
|
|
return cloned
|
|
}
|
|
|
|
func TestFileCooldownStateStore_StateRelativePath(t *testing.T) {
|
|
authDir := filepath.Join(t.TempDir(), "auths")
|
|
store := NewFileCooldownStateStoreWithAuthDir(authDir, authDir)
|
|
|
|
cases := []struct {
|
|
name string
|
|
record CooldownStateRecord
|
|
want string
|
|
}{
|
|
{
|
|
name: "absolute auth file under auth dir",
|
|
record: CooldownStateRecord{
|
|
AuthID: "auth-1",
|
|
AuthFile: filepath.Join(authDir, "nested", "xai.json"),
|
|
},
|
|
want: filepath.Join("nested", "xai.cds"),
|
|
},
|
|
{
|
|
name: "relative auth file",
|
|
record: CooldownStateRecord{
|
|
AuthID: "auth-2",
|
|
AuthFile: filepath.Join("team", "xai.json"),
|
|
},
|
|
want: filepath.Join("team", "xai.cds"),
|
|
},
|
|
{
|
|
name: "absolute auth file outside auth dir",
|
|
record: CooldownStateRecord{
|
|
AuthID: "auth-3",
|
|
AuthFile: filepath.Join(t.TempDir(), "outside.json"),
|
|
},
|
|
want: "outside.cds",
|
|
},
|
|
{
|
|
name: "relative parent escape is rejected",
|
|
record: CooldownStateRecord{
|
|
AuthID: "auth-4",
|
|
AuthFile: filepath.Join("..", "escape.json"),
|
|
},
|
|
want: "",
|
|
},
|
|
{
|
|
name: "auth id fallback",
|
|
record: CooldownStateRecord{
|
|
AuthID: "auth/id 5",
|
|
},
|
|
want: "auth_id_5.cds",
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
if got := store.stateRelativePath(tc.record); got != tc.want {
|
|
t.Fatalf("stateRelativePath() = %q, want %q", got, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFileCooldownStateStore_SaveLoadAndCleanStale(t *testing.T) {
|
|
authDir := t.TempDir()
|
|
store := NewFileCooldownStateStoreWithAuthDir(authDir, authDir)
|
|
ctx := context.Background()
|
|
|
|
stalePath := filepath.Join(authDir, "stale.cds")
|
|
if errWrite := os.WriteFile(stalePath, []byte("{}\n"), 0o600); errWrite != nil {
|
|
t.Fatalf("write stale file: %v", errWrite)
|
|
}
|
|
|
|
nextRetry := time.Now().Add(time.Hour).UTC().Truncate(time.Second)
|
|
updatedAt := time.Now().UTC().Truncate(time.Second)
|
|
record := CooldownStateRecord{
|
|
Provider: "xai",
|
|
AuthID: "auth-1",
|
|
AuthFile: filepath.Join(authDir, "xai.json"),
|
|
Model: "grok-4",
|
|
Status: "cooling",
|
|
NextRetryAfter: nextRetry,
|
|
Reason: "quota",
|
|
Quota: QuotaState{
|
|
Exceeded: true,
|
|
Reason: "quota",
|
|
NextRecoverAt: nextRetry,
|
|
BackoffLevel: 1,
|
|
},
|
|
LastError: &Error{Message: "rate limited", HTTPStatus: 429},
|
|
UpdatedAt: updatedAt,
|
|
}
|
|
|
|
if errSave := store.Save(ctx, []CooldownStateRecord{record}); errSave != nil {
|
|
t.Fatalf("Save() returned error: %v", errSave)
|
|
}
|
|
if _, errStat := os.Stat(filepath.Join(authDir, "xai.cds")); errStat != nil {
|
|
t.Fatalf("expected xai.cds to exist: %v", errStat)
|
|
}
|
|
if _, errStat := os.Stat(stalePath); !errors.Is(errStat, os.ErrNotExist) {
|
|
t.Fatalf("expected stale.cds to be removed, stat error = %v", errStat)
|
|
}
|
|
|
|
loaded, errLoad := store.Load(ctx)
|
|
if errLoad != nil {
|
|
t.Fatalf("Load() returned error: %v", errLoad)
|
|
}
|
|
if len(loaded) != 1 {
|
|
t.Fatalf("loaded records = %d, want 1", len(loaded))
|
|
}
|
|
if loaded[0].AuthID != record.AuthID || loaded[0].Model != record.Model || !loaded[0].NextRetryAfter.Equal(nextRetry) {
|
|
t.Fatalf("loaded record = %+v, want auth/model/retry from %+v", loaded[0], record)
|
|
}
|
|
if loaded[0].LastError == nil || loaded[0].LastError.HTTPStatus != 429 {
|
|
t.Fatalf("loaded last error = %+v, want HTTP 429", loaded[0].LastError)
|
|
}
|
|
|
|
if errSave := store.Save(ctx, nil); errSave != nil {
|
|
t.Fatalf("Save(nil) returned error: %v", errSave)
|
|
}
|
|
if _, errStat := os.Stat(filepath.Join(authDir, "xai.cds")); !errors.Is(errStat, os.ErrNotExist) {
|
|
t.Fatalf("expected xai.cds to be removed, stat error = %v", errStat)
|
|
}
|
|
}
|
|
|
|
func TestFileCooldownStateStore_ConcurrentSave(t *testing.T) {
|
|
authDir := t.TempDir()
|
|
store := NewFileCooldownStateStoreWithAuthDir(authDir, authDir)
|
|
ctx := context.Background()
|
|
nextRetry := time.Now().Add(time.Hour).UTC().Truncate(time.Second)
|
|
|
|
var wg sync.WaitGroup
|
|
errs := make(chan error, 16)
|
|
for i := 0; i < 16; i++ {
|
|
i := i
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
errs <- store.Save(ctx, []CooldownStateRecord{
|
|
{
|
|
Provider: "xai",
|
|
AuthID: "auth-1",
|
|
AuthFile: filepath.Join(authDir, "xai.json"),
|
|
Model: "grok-4",
|
|
Status: "cooling",
|
|
NextRetryAfter: nextRetry.Add(time.Duration(i) * time.Second),
|
|
UpdatedAt: nextRetry,
|
|
},
|
|
})
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
close(errs)
|
|
for errSave := range errs {
|
|
if errSave != nil {
|
|
t.Fatalf("Save() returned error: %v", errSave)
|
|
}
|
|
}
|
|
|
|
loaded, errLoad := store.Load(ctx)
|
|
if errLoad != nil {
|
|
t.Fatalf("Load() returned error: %v", errLoad)
|
|
}
|
|
if len(loaded) != 1 {
|
|
t.Fatalf("loaded records = %d, want 1", len(loaded))
|
|
}
|
|
|
|
tmpMatches, errGlob := filepath.Glob(filepath.Join(authDir, "*.tmp"))
|
|
if errGlob != nil {
|
|
t.Fatalf("glob temp files: %v", errGlob)
|
|
}
|
|
if len(tmpMatches) != 0 {
|
|
t.Fatalf("leftover temp files = %v, want none", tmpMatches)
|
|
}
|
|
}
|
|
|
|
func TestManager_MarkResult_PersistsCooldownOnlyWhenStateChanges(t *testing.T) {
|
|
store := &recordingCooldownStateStore{}
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetCooldownStateStore(store)
|
|
|
|
auth := &Auth{ID: "auth-1", Provider: "xai", Status: StatusActive}
|
|
if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
|
|
t.Fatalf("Register() returned error: %v", errRegister)
|
|
}
|
|
|
|
manager.MarkResult(context.Background(), Result{AuthID: auth.ID, Provider: "xai", Model: "grok-4", Success: true})
|
|
if got := store.saveCount.Load(); got != 0 {
|
|
t.Fatalf("healthy success saved cooldown state %d times, want 0", got)
|
|
}
|
|
|
|
manager.MarkResult(context.Background(), Result{
|
|
AuthID: auth.ID,
|
|
Provider: "xai",
|
|
Model: "grok-4",
|
|
Success: false,
|
|
Error: &Error{Message: "upstream unavailable", HTTPStatus: 500},
|
|
})
|
|
if got := store.saveCount.Load(); got != 1 {
|
|
t.Fatalf("cooldown failure saved cooldown state %d times, want 1", got)
|
|
}
|
|
|
|
manager.MarkResult(context.Background(), Result{AuthID: auth.ID, Provider: "xai", Model: "grok-4", Success: true})
|
|
if got := store.saveCount.Load(); got != 2 {
|
|
t.Fatalf("cooldown clear saved cooldown state %d times, want 2", got)
|
|
}
|
|
|
|
manager.MarkResult(context.Background(), Result{AuthID: auth.ID, Provider: "xai", Model: "grok-4", Success: true})
|
|
if got := store.saveCount.Load(); got != 2 {
|
|
t.Fatalf("clean success saved cooldown state %d times, want 2", got)
|
|
}
|
|
}
|
|
|
|
func TestManager_RestoreCooldownStates(t *testing.T) {
|
|
nextRetry := time.Now().Add(time.Hour).UTC().Truncate(time.Second)
|
|
store := &recordingCooldownStateStore{
|
|
load: []CooldownStateRecord{
|
|
{
|
|
Provider: "xai",
|
|
AuthID: "auth-1",
|
|
Model: "grok-4",
|
|
Status: "cooling",
|
|
NextRetryAfter: nextRetry,
|
|
Reason: "quota",
|
|
Quota: QuotaState{
|
|
Exceeded: true,
|
|
Reason: "quota",
|
|
NextRecoverAt: nextRetry,
|
|
},
|
|
LastError: &Error{Message: "rate limited", HTTPStatus: 429},
|
|
UpdatedAt: nextRetry.Add(-time.Minute),
|
|
},
|
|
},
|
|
}
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetCooldownStateStore(store)
|
|
if _, errRegister := manager.Register(WithSkipPersist(context.Background()), &Auth{ID: "auth-1", Provider: "xai"}); errRegister != nil {
|
|
t.Fatalf("Register() returned error: %v", errRegister)
|
|
}
|
|
|
|
if errRestore := manager.RestoreCooldownStates(context.Background()); errRestore != nil {
|
|
t.Fatalf("RestoreCooldownStates() returned error: %v", errRestore)
|
|
}
|
|
|
|
auth, ok := manager.GetByID("auth-1")
|
|
if !ok {
|
|
t.Fatal("restored auth was not found")
|
|
}
|
|
state := auth.ModelStates["grok-4"]
|
|
if state == nil {
|
|
t.Fatal("model state was not restored")
|
|
}
|
|
if !state.Unavailable || state.Status != StatusError || !state.NextRetryAfter.Equal(nextRetry) {
|
|
t.Fatalf("restored state = %+v, want unavailable status error until %v", state, nextRetry)
|
|
}
|
|
if state.LastError == nil || state.LastError.HTTPStatus != 429 {
|
|
t.Fatalf("restored last error = %+v, want HTTP 429", state.LastError)
|
|
}
|
|
if got := store.saveCount.Load(); got != 1 {
|
|
t.Fatalf("restore cleanup saved cooldown state %d times, want 1", got)
|
|
}
|
|
}
|