Commit e298a718 authored by weak-fox's avatar weak-fox
Browse files

fix: clear temp unsched when resetting account status

parent fdd8499f
......@@ -1866,6 +1866,18 @@ func (s *adminServiceImpl) ClearAccountError(ctx context.Context, id int64) (*Ac
if err := s.accountRepo.ClearError(ctx, id); err != nil {
return nil, err
}
if err := s.accountRepo.ClearRateLimit(ctx, id); err != nil {
return nil, err
}
if err := s.accountRepo.ClearAntigravityQuotaScopes(ctx, id); err != nil {
return nil, err
}
if err := s.accountRepo.ClearModelRateLimits(ctx, id); err != nil {
return nil, err
}
if err := s.accountRepo.ClearTempUnschedulable(ctx, id); err != nil {
return nil, err
}
return s.accountRepo.GetByID(ctx, id)
}
......
//go:build unit
package service
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type accountRepoStubForClearAccountError struct {
mockAccountRepoForGemini
account *Account
clearErrorCalls int
clearRateLimitCalls int
clearAntigravityCalls int
clearModelRateLimitCalls int
clearTempUnschedCalls int
}
func (r *accountRepoStubForClearAccountError) GetByID(ctx context.Context, id int64) (*Account, error) {
return r.account, nil
}
func (r *accountRepoStubForClearAccountError) ClearError(ctx context.Context, id int64) error {
r.clearErrorCalls++
r.account.Status = StatusActive
r.account.ErrorMessage = ""
return nil
}
func (r *accountRepoStubForClearAccountError) ClearRateLimit(ctx context.Context, id int64) error {
r.clearRateLimitCalls++
r.account.RateLimitedAt = nil
r.account.RateLimitResetAt = nil
return nil
}
func (r *accountRepoStubForClearAccountError) ClearAntigravityQuotaScopes(ctx context.Context, id int64) error {
r.clearAntigravityCalls++
return nil
}
func (r *accountRepoStubForClearAccountError) ClearModelRateLimits(ctx context.Context, id int64) error {
r.clearModelRateLimitCalls++
return nil
}
func (r *accountRepoStubForClearAccountError) ClearTempUnschedulable(ctx context.Context, id int64) error {
r.clearTempUnschedCalls++
r.account.TempUnschedulableUntil = nil
r.account.TempUnschedulableReason = ""
return nil
}
func TestAdminService_ClearAccountError_AlsoClearsRecoverableRuntimeState(t *testing.T) {
until := time.Now().Add(10 * time.Minute)
resetAt := time.Now().Add(5 * time.Minute)
repo := &accountRepoStubForClearAccountError{
account: &Account{
ID: 31,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusError,
ErrorMessage: "refresh failed",
RateLimitResetAt: &resetAt,
TempUnschedulableUntil: &until,
TempUnschedulableReason: "missing refresh token",
},
}
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.ClearAccountError(context.Background(), 31)
require.NoError(t, err)
require.NotNil(t, updated)
require.Equal(t, 1, repo.clearErrorCalls)
require.Equal(t, 1, repo.clearRateLimitCalls)
require.Equal(t, 1, repo.clearAntigravityCalls)
require.Equal(t, 1, repo.clearModelRateLimitCalls)
require.Equal(t, 1, repo.clearTempUnschedCalls)
require.Nil(t, updated.RateLimitResetAt)
require.Nil(t, updated.TempUnschedulableUntil)
require.Empty(t, updated.TempUnschedulableReason)
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment