From 0b32f610625d8ab08e3bdd9cedfee9ffc900c54d Mon Sep 17 00:00:00 2001 From: yangjianbo Date: Sun, 22 Feb 2026 17:00:29 +0800 Subject: [PATCH] =?UTF-8?q?fix(ratelimit):=20=E6=B8=85=E9=99=A4=E9=99=90?= =?UTF-8?q?=E6=B5=81=E6=97=B6=E5=90=8C=E6=AD=A5=E6=B8=85=E7=90=86=E4=B8=B4?= =?UTF-8?q?=E6=97=B6=E4=B8=8D=E5=8F=AF=E8=B0=83=E5=BA=A6=E7=8A=B6=E6=80=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ClearRateLimit 增加清理 temp_unschedulable 与缓存\n- 新增 ClearRateLimit 相关单元测试覆盖成功与失败分支 --- backend/internal/service/ratelimit_service.go | 14 +- .../service/ratelimit_service_clear_test.go | 172 ++++++++++++++++++ 2 files changed, 185 insertions(+), 1 deletion(-) create mode 100644 backend/internal/service/ratelimit_service_clear_test.go diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index b1d767fc..fcc7c4a0 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -738,7 +738,19 @@ func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64) if err := s.accountRepo.ClearAntigravityQuotaScopes(ctx, accountID); err != nil { return err } - return s.accountRepo.ClearModelRateLimits(ctx, accountID) + if err := s.accountRepo.ClearModelRateLimits(ctx, accountID); err != nil { + return err + } + // 清除限流时一并清理临时不可调度状态，避免周限/窗口重置后仍被本地临时状态阻断。 + if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil { + return err + } + if s.tempUnschedCache != nil { + if err := s.tempUnschedCache.DeleteTempUnsched(ctx, accountID); err != nil { + slog.Warn("temp_unsched_cache_delete_failed", "account_id", accountID, "error", err) + } + } + return nil } func (s *RateLimitService) ClearTempUnschedulable(ctx context.Context, accountID int64) error { diff --git a/backend/internal/service/ratelimit_service_clear_test.go b/backend/internal/service/ratelimit_service_clear_test.go new file mode 100644 index 00000000..f48151ed --- /dev/null +++ b/backend/internal/service/ratelimit_service_clear_test.go @@ -0,0 +1,172 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type rateLimitClearRepoStub struct { + mockAccountRepoForGemini + clearRateLimitCalls int + clearAntigravityCalls int + clearModelRateLimitCalls int + clearTempUnschedCalls int + clearRateLimitErr error + clearAntigravityErr error + clearModelRateLimitErr error + clearTempUnschedulableErr error +} + +func (r *rateLimitClearRepoStub) ClearRateLimit(ctx context.Context, id int64) error { + r.clearRateLimitCalls++ + return r.clearRateLimitErr +} + +func (r *rateLimitClearRepoStub) ClearAntigravityQuotaScopes(ctx context.Context, id int64) error { + r.clearAntigravityCalls++ + return r.clearAntigravityErr +} + +func (r *rateLimitClearRepoStub) ClearModelRateLimits(ctx context.Context, id int64) error { + r.clearModelRateLimitCalls++ + return r.clearModelRateLimitErr +} + +func (r *rateLimitClearRepoStub) ClearTempUnschedulable(ctx context.Context, id int64) error { + r.clearTempUnschedCalls++ + return r.clearTempUnschedulableErr +} + +type tempUnschedCacheRecorder struct { + deletedIDs []int64 + deleteErr error +} + +func (c *tempUnschedCacheRecorder) SetTempUnsched(ctx context.Context, accountID int64, state *TempUnschedState) error { + return nil +} + +func (c *tempUnschedCacheRecorder) GetTempUnsched(ctx context.Context, accountID int64) (*TempUnschedState, error) { + return nil, nil +} + +func (c *tempUnschedCacheRecorder) DeleteTempUnsched(ctx context.Context, accountID int64) error { + c.deletedIDs = append(c.deletedIDs, accountID) + return c.deleteErr +} + +func TestRateLimitService_ClearRateLimit_AlsoClearsTempUnschedulable(t *testing.T) { + repo := &rateLimitClearRepoStub{} + cache := &tempUnschedCacheRecorder{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, cache) + + err := svc.ClearRateLimit(context.Background(), 42) + require.NoError(t, err) + + 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.Equal(t, []int64{42}, cache.deletedIDs) +} + +func TestRateLimitService_ClearRateLimit_ClearTempUnschedulableFailed(t *testing.T) { + repo := &rateLimitClearRepoStub{ + clearTempUnschedulableErr: errors.New("clear temp unsched failed"), + } + cache := &tempUnschedCacheRecorder{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, cache) + + err := svc.ClearRateLimit(context.Background(), 7) + require.Error(t, err) + + require.Equal(t, 1, repo.clearTempUnschedCalls) + require.Empty(t, cache.deletedIDs) +} + +func TestRateLimitService_ClearRateLimit_ClearRateLimitFailed(t *testing.T) { + repo := &rateLimitClearRepoStub{ + clearRateLimitErr: errors.New("clear rate limit failed"), + } + cache := &tempUnschedCacheRecorder{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, cache) + + err := svc.ClearRateLimit(context.Background(), 11) + require.Error(t, err) + + require.Equal(t, 1, repo.clearRateLimitCalls) + require.Equal(t, 0, repo.clearAntigravityCalls) + require.Equal(t, 0, repo.clearModelRateLimitCalls) + require.Equal(t, 0, repo.clearTempUnschedCalls) + require.Empty(t, cache.deletedIDs) +} + +func TestRateLimitService_ClearRateLimit_ClearAntigravityFailed(t *testing.T) { + repo := &rateLimitClearRepoStub{ + clearAntigravityErr: errors.New("clear antigravity failed"), + } + cache := &tempUnschedCacheRecorder{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, cache) + + err := svc.ClearRateLimit(context.Background(), 12) + require.Error(t, err) + + require.Equal(t, 1, repo.clearRateLimitCalls) + require.Equal(t, 1, repo.clearAntigravityCalls) + require.Equal(t, 0, repo.clearModelRateLimitCalls) + require.Equal(t, 0, repo.clearTempUnschedCalls) + require.Empty(t, cache.deletedIDs) +} + +func TestRateLimitService_ClearRateLimit_ClearModelRateLimitsFailed(t *testing.T) { + repo := &rateLimitClearRepoStub{ + clearModelRateLimitErr: errors.New("clear model rate limits failed"), + } + cache := &tempUnschedCacheRecorder{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, cache) + + err := svc.ClearRateLimit(context.Background(), 13) + require.Error(t, err) + + require.Equal(t, 1, repo.clearRateLimitCalls) + require.Equal(t, 1, repo.clearAntigravityCalls) + require.Equal(t, 1, repo.clearModelRateLimitCalls) + require.Equal(t, 0, repo.clearTempUnschedCalls) + require.Empty(t, cache.deletedIDs) +} + +func TestRateLimitService_ClearRateLimit_CacheDeleteFailedShouldNotFail(t *testing.T) { + repo := &rateLimitClearRepoStub{} + cache := &tempUnschedCacheRecorder{ + deleteErr: errors.New("cache delete failed"), + } + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, cache) + + err := svc.ClearRateLimit(context.Background(), 14) + require.NoError(t, err) + + 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.Equal(t, []int64{14}, cache.deletedIDs) +} + +func TestRateLimitService_ClearRateLimit_WithoutTempUnschedCache(t *testing.T) { + repo := &rateLimitClearRepoStub{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + + err := svc.ClearRateLimit(context.Background(), 15) + require.NoError(t, err) + + require.Equal(t, 1, repo.clearRateLimitCalls) + require.Equal(t, 1, repo.clearAntigravityCalls) + require.Equal(t, 1, repo.clearModelRateLimitCalls) + require.Equal(t, 1, repo.clearTempUnschedCalls) +}