|
| 1 | +//go:build unit |
| 2 | + |
| 3 | +package service |
| 4 | + |
| 5 | +import ( |
| 6 | + "context" |
| 7 | + "encoding/json" |
| 8 | + "testing" |
| 9 | + "time" |
| 10 | + |
| 11 | + "github.com/Wei-Shaw/sub2api/internal/config" |
| 12 | + "github.com/stretchr/testify/require" |
| 13 | +) |
| 14 | + |
| 15 | +// --------------------------------------------------------------------------- |
| 16 | +// errSettingRepo: a SettingRepository that always returns errors on read |
| 17 | +// --------------------------------------------------------------------------- |
| 18 | + |
| 19 | +type errSettingRepo struct { |
| 20 | + mockSettingRepo // embed the existing mock from backup_service_test.go |
| 21 | + readErr error |
| 22 | +} |
| 23 | + |
| 24 | +func (r *errSettingRepo) GetValue(_ context.Context, _ string) (string, error) { |
| 25 | + return "", r.readErr |
| 26 | +} |
| 27 | + |
| 28 | +func (r *errSettingRepo) Get(_ context.Context, _ string) (*Setting, error) { |
| 29 | + return nil, r.readErr |
| 30 | +} |
| 31 | + |
| 32 | +// --------------------------------------------------------------------------- |
| 33 | +// overloadAccountRepoStub: records SetOverloaded calls |
| 34 | +// --------------------------------------------------------------------------- |
| 35 | + |
| 36 | +type overloadAccountRepoStub struct { |
| 37 | + mockAccountRepoForGemini |
| 38 | + overloadCalls int |
| 39 | + lastOverloadID int64 |
| 40 | + lastOverloadEnd time.Time |
| 41 | +} |
| 42 | + |
| 43 | +func (r *overloadAccountRepoStub) SetOverloaded(_ context.Context, id int64, until time.Time) error { |
| 44 | + r.overloadCalls++ |
| 45 | + r.lastOverloadID = id |
| 46 | + r.lastOverloadEnd = until |
| 47 | + return nil |
| 48 | +} |
| 49 | + |
| 50 | +// =========================================================================== |
| 51 | +// SettingService: GetOverloadCooldownSettings |
| 52 | +// =========================================================================== |
| 53 | + |
| 54 | +func TestGetOverloadCooldownSettings_DefaultsWhenNotSet(t *testing.T) { |
| 55 | + repo := newMockSettingRepo() |
| 56 | + svc := NewSettingService(repo, &config.Config{}) |
| 57 | + |
| 58 | + settings, err := svc.GetOverloadCooldownSettings(context.Background()) |
| 59 | + require.NoError(t, err) |
| 60 | + require.True(t, settings.Enabled) |
| 61 | + require.Equal(t, 10, settings.CooldownMinutes) |
| 62 | +} |
| 63 | + |
| 64 | +func TestGetOverloadCooldownSettings_ReadsFromDB(t *testing.T) { |
| 65 | + repo := newMockSettingRepo() |
| 66 | + data, _ := json.Marshal(OverloadCooldownSettings{Enabled: false, CooldownMinutes: 30}) |
| 67 | + repo.data[SettingKeyOverloadCooldownSettings] = string(data) |
| 68 | + svc := NewSettingService(repo, &config.Config{}) |
| 69 | + |
| 70 | + settings, err := svc.GetOverloadCooldownSettings(context.Background()) |
| 71 | + require.NoError(t, err) |
| 72 | + require.False(t, settings.Enabled) |
| 73 | + require.Equal(t, 30, settings.CooldownMinutes) |
| 74 | +} |
| 75 | + |
| 76 | +func TestGetOverloadCooldownSettings_ClampsMinValue(t *testing.T) { |
| 77 | + repo := newMockSettingRepo() |
| 78 | + data, _ := json.Marshal(OverloadCooldownSettings{Enabled: true, CooldownMinutes: 0}) |
| 79 | + repo.data[SettingKeyOverloadCooldownSettings] = string(data) |
| 80 | + svc := NewSettingService(repo, &config.Config{}) |
| 81 | + |
| 82 | + settings, err := svc.GetOverloadCooldownSettings(context.Background()) |
| 83 | + require.NoError(t, err) |
| 84 | + require.Equal(t, 1, settings.CooldownMinutes) |
| 85 | +} |
| 86 | + |
| 87 | +func TestGetOverloadCooldownSettings_ClampsMaxValue(t *testing.T) { |
| 88 | + repo := newMockSettingRepo() |
| 89 | + data, _ := json.Marshal(OverloadCooldownSettings{Enabled: true, CooldownMinutes: 999}) |
| 90 | + repo.data[SettingKeyOverloadCooldownSettings] = string(data) |
| 91 | + svc := NewSettingService(repo, &config.Config{}) |
| 92 | + |
| 93 | + settings, err := svc.GetOverloadCooldownSettings(context.Background()) |
| 94 | + require.NoError(t, err) |
| 95 | + require.Equal(t, 120, settings.CooldownMinutes) |
| 96 | +} |
| 97 | + |
| 98 | +func TestGetOverloadCooldownSettings_InvalidJSON_ReturnsDefaults(t *testing.T) { |
| 99 | + repo := newMockSettingRepo() |
| 100 | + repo.data[SettingKeyOverloadCooldownSettings] = "not-json" |
| 101 | + svc := NewSettingService(repo, &config.Config{}) |
| 102 | + |
| 103 | + settings, err := svc.GetOverloadCooldownSettings(context.Background()) |
| 104 | + require.NoError(t, err) |
| 105 | + require.True(t, settings.Enabled) |
| 106 | + require.Equal(t, 10, settings.CooldownMinutes) |
| 107 | +} |
| 108 | + |
| 109 | +func TestGetOverloadCooldownSettings_EmptyValue_ReturnsDefaults(t *testing.T) { |
| 110 | + repo := newMockSettingRepo() |
| 111 | + repo.data[SettingKeyOverloadCooldownSettings] = "" |
| 112 | + svc := NewSettingService(repo, &config.Config{}) |
| 113 | + |
| 114 | + settings, err := svc.GetOverloadCooldownSettings(context.Background()) |
| 115 | + require.NoError(t, err) |
| 116 | + require.True(t, settings.Enabled) |
| 117 | + require.Equal(t, 10, settings.CooldownMinutes) |
| 118 | +} |
| 119 | + |
| 120 | +// =========================================================================== |
| 121 | +// SettingService: SetOverloadCooldownSettings |
| 122 | +// =========================================================================== |
| 123 | + |
| 124 | +func TestSetOverloadCooldownSettings_Success(t *testing.T) { |
| 125 | + repo := newMockSettingRepo() |
| 126 | + svc := NewSettingService(repo, &config.Config{}) |
| 127 | + |
| 128 | + err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{ |
| 129 | + Enabled: false, |
| 130 | + CooldownMinutes: 25, |
| 131 | + }) |
| 132 | + require.NoError(t, err) |
| 133 | + |
| 134 | + // Verify round-trip |
| 135 | + settings, err := svc.GetOverloadCooldownSettings(context.Background()) |
| 136 | + require.NoError(t, err) |
| 137 | + require.False(t, settings.Enabled) |
| 138 | + require.Equal(t, 25, settings.CooldownMinutes) |
| 139 | +} |
| 140 | + |
| 141 | +func TestSetOverloadCooldownSettings_RejectsNil(t *testing.T) { |
| 142 | + svc := NewSettingService(newMockSettingRepo(), &config.Config{}) |
| 143 | + err := svc.SetOverloadCooldownSettings(context.Background(), nil) |
| 144 | + require.Error(t, err) |
| 145 | +} |
| 146 | + |
| 147 | +func TestSetOverloadCooldownSettings_EnabledRejectsOutOfRange(t *testing.T) { |
| 148 | + svc := NewSettingService(newMockSettingRepo(), &config.Config{}) |
| 149 | + |
| 150 | + for _, minutes := range []int{0, -1, 121, 999} { |
| 151 | + err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{ |
| 152 | + Enabled: true, CooldownMinutes: minutes, |
| 153 | + }) |
| 154 | + require.Error(t, err, "should reject enabled=true + cooldown_minutes=%d", minutes) |
| 155 | + require.Contains(t, err.Error(), "cooldown_minutes must be between 1-120") |
| 156 | + } |
| 157 | +} |
| 158 | + |
| 159 | +func TestSetOverloadCooldownSettings_DisabledNormalizesOutOfRange(t *testing.T) { |
| 160 | + repo := newMockSettingRepo() |
| 161 | + svc := NewSettingService(repo, &config.Config{}) |
| 162 | + |
| 163 | + // enabled=false + cooldown_minutes=0 应该保存成功,值被归一化为10 |
| 164 | + err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{ |
| 165 | + Enabled: false, CooldownMinutes: 0, |
| 166 | + }) |
| 167 | + require.NoError(t, err, "disabled with invalid minutes should NOT be rejected") |
| 168 | + |
| 169 | + // 验证持久化后读回来的值 |
| 170 | + settings, err := svc.GetOverloadCooldownSettings(context.Background()) |
| 171 | + require.NoError(t, err) |
| 172 | + require.False(t, settings.Enabled) |
| 173 | + require.Equal(t, 10, settings.CooldownMinutes, "should be normalized to default") |
| 174 | +} |
| 175 | + |
| 176 | +func TestSetOverloadCooldownSettings_AcceptsBoundaries(t *testing.T) { |
| 177 | + svc := NewSettingService(newMockSettingRepo(), &config.Config{}) |
| 178 | + |
| 179 | + for _, minutes := range []int{1, 60, 120} { |
| 180 | + err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{ |
| 181 | + Enabled: true, CooldownMinutes: minutes, |
| 182 | + }) |
| 183 | + require.NoError(t, err, "should accept cooldown_minutes=%d", minutes) |
| 184 | + } |
| 185 | +} |
| 186 | + |
| 187 | +// =========================================================================== |
| 188 | +// RateLimitService: handle529 behaviour |
| 189 | +// =========================================================================== |
| 190 | + |
| 191 | +func TestHandle529_EnabledFromDB_PausesAccount(t *testing.T) { |
| 192 | + accountRepo := &overloadAccountRepoStub{} |
| 193 | + settingRepo := newMockSettingRepo() |
| 194 | + data, _ := json.Marshal(OverloadCooldownSettings{Enabled: true, CooldownMinutes: 15}) |
| 195 | + settingRepo.data[SettingKeyOverloadCooldownSettings] = string(data) |
| 196 | + |
| 197 | + settingSvc := NewSettingService(settingRepo, &config.Config{}) |
| 198 | + svc := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil) |
| 199 | + svc.SetSettingService(settingSvc) |
| 200 | + |
| 201 | + account := &Account{ID: 42, Platform: PlatformAnthropic, Type: AccountTypeOAuth} |
| 202 | + before := time.Now() |
| 203 | + svc.handle529(context.Background(), account) |
| 204 | + |
| 205 | + require.Equal(t, 1, accountRepo.overloadCalls) |
| 206 | + require.Equal(t, int64(42), accountRepo.lastOverloadID) |
| 207 | + require.WithinDuration(t, before.Add(15*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second) |
| 208 | +} |
| 209 | + |
| 210 | +func TestHandle529_DisabledFromDB_SkipsAccount(t *testing.T) { |
| 211 | + accountRepo := &overloadAccountRepoStub{} |
| 212 | + settingRepo := newMockSettingRepo() |
| 213 | + data, _ := json.Marshal(OverloadCooldownSettings{Enabled: false, CooldownMinutes: 15}) |
| 214 | + settingRepo.data[SettingKeyOverloadCooldownSettings] = string(data) |
| 215 | + |
| 216 | + settingSvc := NewSettingService(settingRepo, &config.Config{}) |
| 217 | + svc := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil) |
| 218 | + svc.SetSettingService(settingSvc) |
| 219 | + |
| 220 | + account := &Account{ID: 42, Platform: PlatformAnthropic, Type: AccountTypeOAuth} |
| 221 | + svc.handle529(context.Background(), account) |
| 222 | + |
| 223 | + require.Equal(t, 0, accountRepo.overloadCalls, "should NOT pause when disabled") |
| 224 | +} |
| 225 | + |
| 226 | +func TestHandle529_NilSettingService_FallsBackToConfig(t *testing.T) { |
| 227 | + accountRepo := &overloadAccountRepoStub{} |
| 228 | + cfg := &config.Config{} |
| 229 | + cfg.RateLimit.OverloadCooldownMinutes = 20 |
| 230 | + svc := NewRateLimitService(accountRepo, nil, cfg, nil, nil) |
| 231 | + // NOT calling SetSettingService — remains nil |
| 232 | + |
| 233 | + account := &Account{ID: 77, Platform: PlatformAnthropic, Type: AccountTypeOAuth} |
| 234 | + before := time.Now() |
| 235 | + svc.handle529(context.Background(), account) |
| 236 | + |
| 237 | + require.Equal(t, 1, accountRepo.overloadCalls) |
| 238 | + require.WithinDuration(t, before.Add(20*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second) |
| 239 | +} |
| 240 | + |
| 241 | +func TestHandle529_NilSettingService_ZeroConfig_DefaultsTen(t *testing.T) { |
| 242 | + accountRepo := &overloadAccountRepoStub{} |
| 243 | + svc := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil) |
| 244 | + |
| 245 | + account := &Account{ID: 88, Platform: PlatformAnthropic, Type: AccountTypeOAuth} |
| 246 | + before := time.Now() |
| 247 | + svc.handle529(context.Background(), account) |
| 248 | + |
| 249 | + require.Equal(t, 1, accountRepo.overloadCalls) |
| 250 | + require.WithinDuration(t, before.Add(10*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second) |
| 251 | +} |
| 252 | + |
| 253 | +func TestHandle529_DBReadError_FallsBackToConfig(t *testing.T) { |
| 254 | + accountRepo := &overloadAccountRepoStub{} |
| 255 | + errRepo := &errSettingRepo{readErr: context.DeadlineExceeded} |
| 256 | + errRepo.data = make(map[string]string) |
| 257 | + |
| 258 | + cfg := &config.Config{} |
| 259 | + cfg.RateLimit.OverloadCooldownMinutes = 7 |
| 260 | + settingSvc := NewSettingService(errRepo, cfg) |
| 261 | + svc := NewRateLimitService(accountRepo, nil, cfg, nil, nil) |
| 262 | + svc.SetSettingService(settingSvc) |
| 263 | + |
| 264 | + account := &Account{ID: 99, Platform: PlatformAnthropic, Type: AccountTypeOAuth} |
| 265 | + before := time.Now() |
| 266 | + svc.handle529(context.Background(), account) |
| 267 | + |
| 268 | + require.Equal(t, 1, accountRepo.overloadCalls) |
| 269 | + require.WithinDuration(t, before.Add(7*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second) |
| 270 | +} |
| 271 | + |
| 272 | +// =========================================================================== |
| 273 | +// Model: defaults & JSON round-trip |
| 274 | +// =========================================================================== |
| 275 | + |
| 276 | +func TestDefaultOverloadCooldownSettings(t *testing.T) { |
| 277 | + d := DefaultOverloadCooldownSettings() |
| 278 | + require.True(t, d.Enabled) |
| 279 | + require.Equal(t, 10, d.CooldownMinutes) |
| 280 | +} |
| 281 | + |
| 282 | +func TestOverloadCooldownSettings_JSONRoundTrip(t *testing.T) { |
| 283 | + original := OverloadCooldownSettings{Enabled: false, CooldownMinutes: 42} |
| 284 | + data, err := json.Marshal(original) |
| 285 | + require.NoError(t, err) |
| 286 | + |
| 287 | + var decoded OverloadCooldownSettings |
| 288 | + require.NoError(t, json.Unmarshal(data, &decoded)) |
| 289 | + require.Equal(t, original, decoded) |
| 290 | + |
| 291 | + // Verify JSON uses snake_case field names |
| 292 | + var raw map[string]any |
| 293 | + require.NoError(t, json.Unmarshal(data, &raw)) |
| 294 | + _, hasEnabled := raw["enabled"] |
| 295 | + _, hasCooldown := raw["cooldown_minutes"] |
| 296 | + require.True(t, hasEnabled, "JSON must use 'enabled'") |
| 297 | + require.True(t, hasCooldown, "JSON must use 'cooldown_minutes'") |
| 298 | +} |
0 commit comments