fix(ratelimit): 清除限流时同步清理临时不可调度状态
- ClearRateLimit 增加清理 temp_unschedulable 与缓存\n- 新增 ClearRateLimit 相关单元测试覆盖成功与失败分支
This commit is contained in:
@@ -738,7 +738,19 @@ func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64)
|
|||||||
if err := s.accountRepo.ClearAntigravityQuotaScopes(ctx, accountID); err != nil {
|
if err := s.accountRepo.ClearAntigravityQuotaScopes(ctx, accountID); err != nil {
|
||||||
return err
|
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 {
|
func (s *RateLimitService) ClearTempUnschedulable(ctx context.Context, accountID int64) error {
|
||||||
|
|||||||
172
backend/internal/service/ratelimit_service_clear_test.go
Normal file
172
backend/internal/service/ratelimit_service_clear_test.go
Normal file
@@ -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)
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user