Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
9.6 kB · 297 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298package store
import ( "context" "database/sql" "database/sql/driver" "errors" "fmt" "io" "reflect" "strings" "sync" "sync/atomic" "testing" "time"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth")
var cooldownTestDriverID atomic.Uint64
type cooldownTestDriver struct { state *cooldownTestState}
type cooldownTestState struct { mu sync.Mutex rows map[string]cooldownTestRow queries []string}
type cooldownTestRow struct { content []byte deleted bool updatedAt time.Time}
type cooldownTestConn struct { state *cooldownTestState}
type cooldownTestTx struct{}
type cooldownTestRows struct { rows []cooldownTestRow index int}
func (d *cooldownTestDriver) Open(string) (driver.Conn, error) { return &cooldownTestConn{state: d.state}, nil}
func (c *cooldownTestConn) Prepare(string) (driver.Stmt, error) { return nil, errors.New("prepare is not supported")}
func (c *cooldownTestConn) Close() error { return nil}
func (c *cooldownTestConn) Begin() (driver.Tx, error) { return &cooldownTestTx{}, nil}
func (c *cooldownTestConn) ExecContext(_ context.Context, query string, args []driver.NamedValue) (driver.Result, error) { c.state.mu.Lock() defer c.state.mu.Unlock() c.state.queries = append(c.state.queries, query) if !strings.Contains(query, "INSERT INTO") || (len(args) != 4 && len(args) != 5) { return driver.RowsAffected(1), nil } authID, okAuthID := args[0].Value.(string) model, okModel := args[1].Value.(string) content, okContent := args[2].Value.([]byte) updatedAt, okUpdatedAt := args[3].Value.(time.Time) if !okAuthID || !okModel || !okContent || !okUpdatedAt { return nil, errors.New("invalid cooldown query arguments") } key := authID + "\x00" + model current, exists := c.state.rows[key] if len(args) == 4 { if !exists || !current.updatedAt.After(updatedAt) { c.state.rows[key] = cooldownTestRow{content: append([]byte(nil), content...), updatedAt: updatedAt} } return driver.RowsAffected(1), nil } observedAt, okObservedAt := args[4].Value.(time.Time) if !okObservedAt { return nil, errors.New("invalid cooldown delete version") } if !exists || (!current.deleted && !current.updatedAt.After(observedAt)) { c.state.rows[key] = cooldownTestRow{content: append([]byte(nil), content...), deleted: true, updatedAt: updatedAt} } return driver.RowsAffected(1), nil}
func (c *cooldownTestConn) QueryContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Rows, error) { c.state.mu.Lock() defer c.state.mu.Unlock() c.state.queries = append(c.state.queries, query) rows := make([]cooldownTestRow, 0, len(c.state.rows)) for _, row := range c.state.rows { if !row.deleted { row.content = append([]byte(nil), row.content...) rows = append(rows, row) } } return &cooldownTestRows{rows: rows}, nil}
func (*cooldownTestTx) Commit() error { return nil}
func (*cooldownTestTx) Rollback() error { return nil}
func (r *cooldownTestRows) Columns() []string { return []string{"content", "updated_at"}}
func (r *cooldownTestRows) Close() error { return nil}
func (r *cooldownTestRows) Next(dest []driver.Value) error { if r.index >= len(r.rows) { return io.EOF } dest[0] = r.rows[r.index].content dest[1] = r.rows[r.index].updatedAt r.index++ return nil}
func TestPostgresCooldownStateStore_SaveLoad(t *testing.T) { state := &cooldownTestState{rows: make(map[string]cooldownTestRow)} driverName := fmt.Sprintf("cliproxy_postgres_cooldown_test_%d", cooldownTestDriverID.Add(1)) sql.Register(driverName, &cooldownTestDriver{state: state}) db, errOpen := sql.Open(driverName, "") if errOpen != nil { t.Fatalf("sql.Open() error = %v", errOpen) } t.Cleanup(func() { if errClose := db.Close(); errClose != nil { t.Errorf("db.Close() error = %v", errClose) } })
postgresStore := &PostgresStore{ db: db, cfg: PostgresStoreConfig{ ConfigTable: defaultConfigTable, AuthTable: defaultAuthTable, CooldownTable: defaultCooldownTable, }, } cooldownStore := &postgresCooldownStateStore{store: postgresStore} postgresStore.cooldownStore = cooldownStore
if errSchema := postgresStore.EnsureSchema(context.Background()); errSchema != nil { t.Fatalf("EnsureSchema() error = %v", errSchema) } if got := postgresStore.CooldownStateStore(); got != cooldownStore { t.Fatalf("CooldownStateStore() = %T, want configured PostgreSQL store", got) }
nextRetry := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) records := []cliproxyauth.CooldownStateRecord{ { Provider: "codex", AuthID: "account-1", Model: "gpt-test", Status: string(cliproxyauth.StatusError), NextRetryAfter: nextRetry, Reason: "rate limited", UpdatedAt: nextRetry.Add(-time.Minute), }, } if errSave := cooldownStore.Save(context.Background(), records); errSave != nil { t.Fatalf("Save() error = %v", errSave) } loaded, errLoad := cooldownStore.Load(context.Background()) if errLoad != nil { t.Fatalf("Load() error = %v", errLoad) } if !reflect.DeepEqual(loaded, records) { t.Fatalf("Load() = %#v, want %#v", loaded, records) }
zeroTimeRecord := cliproxyauth.CooldownStateRecord{AuthID: "account-2", Model: "gpt-test"} if errSave := cooldownStore.Save(context.Background(), []cliproxyauth.CooldownStateRecord{zeroTimeRecord}); errSave != nil { t.Fatalf("Save() with zero UpdatedAt error = %v", errSave) } loaded, errLoad = cooldownStore.Load(context.Background()) if errLoad != nil { t.Fatalf("Load() after zero UpdatedAt error = %v", errLoad) } if len(loaded) != 1 || loaded[0].UpdatedAt.IsZero() { t.Fatalf("Load() did not persist a normalized UpdatedAt: %#v", loaded) }
if errSave := cooldownStore.Save(context.Background(), nil); errSave != nil { t.Fatalf("Save(nil) error = %v", errSave) } loaded, errLoad = cooldownStore.Load(context.Background()) if errLoad != nil { t.Fatalf("Load() after Save(nil) error = %v", errLoad) } if len(loaded) != 0 { t.Fatalf("Load() after Save(nil) returned %d records, want 0", len(loaded)) }
state.mu.Lock() queries := strings.Join(state.queries, "\n") state.mu.Unlock() if !strings.Contains(queries, `CREATE TABLE IF NOT EXISTS "cooldown_store"`) { t.Fatalf("EnsureSchema() did not create cooldown table; queries:\n%s", queries) }}
func TestPostgresCooldownStateStore_MergesConcurrentInstances(t *testing.T) { state := &cooldownTestState{rows: make(map[string]cooldownTestRow)} driverName := fmt.Sprintf("cliproxy_postgres_cooldown_merge_test_%d", cooldownTestDriverID.Add(1)) sql.Register(driverName, &cooldownTestDriver{state: state}) db, errOpen := sql.Open(driverName, "") if errOpen != nil { t.Fatalf("sql.Open() error = %v", errOpen) } t.Cleanup(func() { if errClose := db.Close(); errClose != nil { t.Errorf("db.Close() error = %v", errClose) } }) postgresStore := &PostgresStore{ db: db, cfg: PostgresStoreConfig{CooldownTable: defaultCooldownTable}, } storeA := &postgresCooldownStateStore{store: postgresStore} storeB := &postgresCooldownStateStore{store: postgresStore} staleStore := &postgresCooldownStateStore{store: postgresStore}
for _, cooldownStore := range []*postgresCooldownStateStore{storeA, storeB} { if _, errLoad := cooldownStore.Load(context.Background()); errLoad != nil { t.Fatalf("initial Load() error = %v", errLoad) } } updatedAt := time.Now().UTC().Add(-time.Minute) recordA := cliproxyauth.CooldownStateRecord{AuthID: "account-a", Model: "model-a", UpdatedAt: updatedAt} recordB := cliproxyauth.CooldownStateRecord{AuthID: "account-b", Model: "model-b", UpdatedAt: updatedAt} if errSave := storeA.Save(context.Background(), []cliproxyauth.CooldownStateRecord{recordA}); errSave != nil { t.Fatalf("storeA.Save() error = %v", errSave) } if errSave := storeB.Save(context.Background(), []cliproxyauth.CooldownStateRecord{recordB}); errSave != nil { t.Fatalf("storeB.Save() error = %v", errSave) } staleRecords, errLoad := staleStore.Load(context.Background()) if errLoad != nil { t.Fatalf("staleStore.Load() error = %v", errLoad) } if len(staleRecords) != 2 { t.Fatalf("merged Load() returned %d records, want 2", len(staleRecords)) }
newerRecordA := recordA newerRecordA.UpdatedAt = updatedAt.Add(time.Hour) if errSave := storeA.Save(context.Background(), []cliproxyauth.CooldownStateRecord{newerRecordA}); errSave != nil { t.Fatalf("storeA.Save(newer) error = %v", errSave) } if errSave := staleStore.Save(context.Background(), []cliproxyauth.CooldownStateRecord{recordB}); errSave != nil { t.Fatalf("staleStore.Save(without newer record) error = %v", errSave) } resurrectStore := &postgresCooldownStateStore{store: postgresStore} activeRecords, errLoad := resurrectStore.Load(context.Background()) if errLoad != nil { t.Fatalf("resurrectStore.Load() error = %v", errLoad) } if len(activeRecords) != 2 { t.Fatalf("Load() after stale delete returned %d records, want 2", len(activeRecords)) }
if errSave := storeA.Save(context.Background(), nil); errSave != nil { t.Fatalf("storeA.Save(nil) error = %v", errSave) } if errSave := resurrectStore.Save(context.Background(), activeRecords); errSave != nil { t.Fatalf("resurrectStore.Save() error = %v", errSave) } reader := &postgresCooldownStateStore{store: postgresStore} loaded, errLoad := reader.Load(context.Background()) if errLoad != nil { t.Fatalf("reader.Load() error = %v", errLoad) } if len(loaded) != 1 || loaded[0].AuthID != recordB.AuthID { t.Fatalf("Load() after stale save = %#v, want only account-b", loaded) }}