ec6bcfeb83
OAuth 账号收到 401 时，原逻辑同时设置 expires_at=now() 和 SetError()， 但刷新服务只查询 status=active 的账号，导致 error 状态的账号永远无法 被刷新服务拾取，expires_at=now() 实际上是死代码。 修复: - OAuth 401 使用 SetTempUnschedulable 替代 SetError，保持 status=active - 新增 oauth_401_cooldown_minutes 配置项（默认 10 分钟） - 刷新成功后同步清除 DB 和 Redis 中的临时不可调度状态 - 不可重试错误检查(invalid_grant 等)从 Antigravity 推广到所有平台 - 可重试错误耗尽后不再标记 error，下个刷新周期继续重试 恢复流程: OAuth 401 → temp_unschedulable + expires_at=now → 刷新服务拾取 → 成功: 清除 temp_unschedulable → 自动恢复 → invalid_grant: SetError → 永久禁用 → 网络错误: 仅记日志 → 下周期重试
122 lines
3.4 KiB
Go
122 lines
3.4 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net/http"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type rateLimitAccountRepoStub struct {
|
|
mockAccountRepoForGemini
|
|
setErrorCalls int
|
|
tempCalls int
|
|
lastErrorMsg string
|
|
}
|
|
|
|
func (r *rateLimitAccountRepoStub) SetError(ctx context.Context, id int64, errorMsg string) error {
|
|
r.setErrorCalls++
|
|
r.lastErrorMsg = errorMsg
|
|
return nil
|
|
}
|
|
|
|
func (r *rateLimitAccountRepoStub) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
|
|
r.tempCalls++
|
|
return nil
|
|
}
|
|
|
|
type tokenCacheInvalidatorRecorder struct {
|
|
accounts []*Account
|
|
err error
|
|
}
|
|
|
|
func (r *tokenCacheInvalidatorRecorder) InvalidateToken(ctx context.Context, account *Account) error {
|
|
r.accounts = append(r.accounts, account)
|
|
return r.err
|
|
}
|
|
|
|
func TestRateLimitService_HandleUpstreamError_OAuth401SetsTempUnschedulable(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
platform string
|
|
}{
|
|
{name: "gemini", platform: PlatformGemini},
|
|
{name: "antigravity", platform: PlatformAntigravity},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
repo := &rateLimitAccountRepoStub{}
|
|
invalidator := &tokenCacheInvalidatorRecorder{}
|
|
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
|
service.SetTokenCacheInvalidator(invalidator)
|
|
account := &Account{
|
|
ID: 100,
|
|
Platform: tt.platform,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"temp_unschedulable_enabled": true,
|
|
"temp_unschedulable_rules": []any{
|
|
map[string]any{
|
|
"error_code": 401,
|
|
"keywords": []any{"unauthorized"},
|
|
"duration_minutes": 30,
|
|
"description": "custom rule",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
shouldDisable := service.HandleUpstreamError(context.Background(), account, 401, http.Header{}, []byte("unauthorized"))
|
|
|
|
require.True(t, shouldDisable)
|
|
require.Equal(t, 0, repo.setErrorCalls)
|
|
require.Equal(t, 1, repo.tempCalls)
|
|
require.Len(t, invalidator.accounts, 1)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRateLimitService_HandleUpstreamError_OAuth401InvalidatorError(t *testing.T) {
|
|
repo := &rateLimitAccountRepoStub{}
|
|
invalidator := &tokenCacheInvalidatorRecorder{err: errors.New("boom")}
|
|
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
|
service.SetTokenCacheInvalidator(invalidator)
|
|
account := &Account{
|
|
ID: 101,
|
|
Platform: PlatformGemini,
|
|
Type: AccountTypeOAuth,
|
|
}
|
|
|
|
shouldDisable := service.HandleUpstreamError(context.Background(), account, 401, http.Header{}, []byte("unauthorized"))
|
|
|
|
require.True(t, shouldDisable)
|
|
require.Equal(t, 0, repo.setErrorCalls)
|
|
require.Equal(t, 1, repo.tempCalls)
|
|
require.Len(t, invalidator.accounts, 1)
|
|
}
|
|
|
|
func TestRateLimitService_HandleUpstreamError_NonOAuth401(t *testing.T) {
|
|
repo := &rateLimitAccountRepoStub{}
|
|
invalidator := &tokenCacheInvalidatorRecorder{}
|
|
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
|
service.SetTokenCacheInvalidator(invalidator)
|
|
account := &Account{
|
|
ID: 102,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeAPIKey,
|
|
}
|
|
|
|
shouldDisable := service.HandleUpstreamError(context.Background(), account, 401, http.Header{}, []byte("unauthorized"))
|
|
|
|
require.True(t, shouldDisable)
|
|
require.Equal(t, 1, repo.setErrorCalls)
|
|
require.Empty(t, invalidator.accounts)
|
|
}
|