ratelimit_service_scheduling_threshold_test.go 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "encoding/json"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/config"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestRateLimitService_ApplyAccountSchedulingThreshold_SetsTempUnschedulable(t *testing.T) {
  12. accountSchedulingThresholdsSF.Forget(SettingKeyAccountSchedulingThresholds)
  13. accountSchedulingThresholdsCache.Store(&cachedAccountSchedulingThresholds{})
  14. settingsRepo := newMockSettingRepo()
  15. settingsRepo.data[SettingKeyAccountSchedulingThresholds] = `{"openai":80}`
  16. accountRepo := &rateLimitAccountRepoStub{}
  17. rl := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
  18. rl.SetSettingService(NewSettingService(settingsRepo, &config.Config{}))
  19. until := time.Now().UTC().Add(6 * time.Hour)
  20. account := &Account{
  21. ID: 1001,
  22. Platform: PlatformOpenAI,
  23. Status: StatusActive,
  24. Schedulable: true,
  25. Extra: map[string]any{
  26. "codex_7d_used_percent": 91.5,
  27. "codex_7d_reset_at": until.Format(time.RFC3339),
  28. },
  29. }
  30. blocked := rl.ApplyAccountSchedulingThreshold(context.Background(), account)
  31. require.True(t, blocked)
  32. require.Equal(t, 1, accountRepo.tempCalls)
  33. require.NotNil(t, account.TempUnschedulableUntil)
  34. require.WithinDuration(t, until, *account.TempUnschedulableUntil, time.Second)
  35. require.True(t, IsAccountSchedulingThresholdReason(accountRepo.lastTempReason))
  36. var payload map[string]any
  37. require.NoError(t, json.Unmarshal([]byte(accountRepo.lastTempReason), &payload))
  38. require.Equal(t, PlatformOpenAI, payload["platform"])
  39. require.Equal(t, "7d", payload["window"])
  40. require.Equal(t, float64(80), payload["threshold_percent"])
  41. require.Equal(t, float64(91.5), payload["used_percent"])
  42. require.Contains(t, payload["error_message"], "91.5% used >= 80%")
  43. }
  44. func TestRateLimitService_ApplyAccountSchedulingThreshold_UsesAccountOverrideInReason(t *testing.T) {
  45. accountSchedulingThresholdsSF.Forget(SettingKeyAccountSchedulingThresholds)
  46. accountSchedulingThresholdsCache.Store(&cachedAccountSchedulingThresholds{})
  47. settingsRepo := newMockSettingRepo()
  48. settingsRepo.data[SettingKeyAccountSchedulingThresholds] = `{"openai":90}`
  49. accountRepo := &rateLimitAccountRepoStub{}
  50. rl := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
  51. rl.SetSettingService(NewSettingService(settingsRepo, &config.Config{}))
  52. until := time.Now().UTC().Add(6 * time.Hour)
  53. account := &Account{
  54. ID: 1003,
  55. Platform: PlatformOpenAI,
  56. Status: StatusActive,
  57. Schedulable: true,
  58. Credentials: map[string]any{
  59. "account_scheduling_threshold": 80,
  60. },
  61. Extra: map[string]any{
  62. "codex_7d_used_percent": 85.5,
  63. "codex_7d_reset_at": until.Format(time.RFC3339),
  64. },
  65. }
  66. blocked := rl.ApplyAccountSchedulingThreshold(context.Background(), account)
  67. require.True(t, blocked)
  68. require.Equal(t, 1, accountRepo.tempCalls)
  69. var payload map[string]any
  70. require.NoError(t, json.Unmarshal([]byte(accountRepo.lastTempReason), &payload))
  71. require.Equal(t, float64(80), payload["threshold_percent"])
  72. require.Equal(t, float64(85.5), payload["used_percent"])
  73. require.Contains(t, payload["error_message"], "85.5% used >= 80%")
  74. }
  75. type fableSchedulingThresholdRepoStub struct {
  76. rateLimitAccountRepoStub
  77. modelCalls int
  78. lastModelScope string
  79. lastModelReset time.Time
  80. lastModelReason string
  81. }
  82. func (r *fableSchedulingThresholdRepoStub) SetModelRateLimit(_ context.Context, _ int64, scope string, resetAt time.Time, reason ...string) error {
  83. r.modelCalls++
  84. r.lastModelScope = scope
  85. r.lastModelReset = resetAt
  86. if len(reason) > 0 {
  87. r.lastModelReason = reason[0]
  88. }
  89. return nil
  90. }
  91. func TestRateLimitService_ApplyAccountSchedulingThreshold_FableOnlyLimitsFableModels(t *testing.T) {
  92. accountSchedulingThresholdsSF.Forget(SettingKeyAccountSchedulingThresholds)
  93. accountSchedulingThresholdsCache.Store(&cachedAccountSchedulingThresholds{})
  94. settingsRepo := newMockSettingRepo()
  95. settingsRepo.data[SettingKeyAccountSchedulingThresholds] = `{"anthropic":100}`
  96. accountRepo := &fableSchedulingThresholdRepoStub{}
  97. rl := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
  98. rl.SetSettingService(NewSettingService(settingsRepo, &config.Config{}))
  99. until := time.Now().UTC().Add(4 * 24 * time.Hour).Truncate(time.Second)
  100. account := &Account{
  101. ID: 1004,
  102. Platform: PlatformAnthropic,
  103. Status: StatusActive,
  104. Schedulable: true,
  105. Credentials: map[string]any{
  106. "account_scheduling_threshold": 60,
  107. },
  108. Extra: map[string]any{
  109. "passive_usage_7d_utilization": 0.40,
  110. "passive_usage_7d_reset": float64(time.Now().UTC().Add(3 * 24 * time.Hour).Unix()),
  111. "passive_usage_7d_oi_utilization": 0.61,
  112. "passive_usage_7d_oi_reset": float64(until.Unix()),
  113. },
  114. }
  115. blocked := rl.ApplyAccountSchedulingThreshold(context.Background(), account)
  116. require.False(t, blocked, "the Fable-only window must not pause the whole account")
  117. require.Zero(t, accountRepo.tempCalls)
  118. require.Equal(t, 1, accountRepo.modelCalls)
  119. require.Equal(t, anthropicFableRateLimitKey, accountRepo.lastModelScope)
  120. require.WithinDuration(t, until, accountRepo.lastModelReset, time.Second)
  121. require.True(t, IsAccountSchedulingThresholdReason(accountRepo.lastModelReason))
  122. require.False(t, account.IsSchedulableForModel("claude-fable-5"))
  123. require.False(t, account.IsSchedulableForModel("claude-fable-5[1m]"))
  124. require.True(t, account.IsSchedulableForModel("claude-opus-4-8"))
  125. require.True(t, account.IsSchedulableForModel("claude-sonnet-4-6"))
  126. blocked = rl.ApplyAccountSchedulingThreshold(context.Background(), account)
  127. require.False(t, blocked)
  128. require.Equal(t, 1, accountRepo.modelCalls, "an active model limit should not be persisted twice")
  129. }
  130. func TestRateLimitService_ApplyAccountSchedulingThreshold_SkipsDuplicateTempUnschedulable(t *testing.T) {
  131. accountSchedulingThresholdsSF.Forget(SettingKeyAccountSchedulingThresholds)
  132. accountSchedulingThresholdsCache.Store(&cachedAccountSchedulingThresholds{})
  133. settingsRepo := newMockSettingRepo()
  134. settingsRepo.data[SettingKeyAccountSchedulingThresholds] = `{"openai":80}`
  135. accountRepo := &rateLimitAccountRepoStub{}
  136. rl := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
  137. rl.SetSettingService(NewSettingService(settingsRepo, &config.Config{}))
  138. until := time.Now().UTC().Add(6 * time.Hour).Truncate(time.Second)
  139. existingReason := BuildDetailedAccountSchedulingThresholdReason(AccountSchedulingThresholdReasonInput{
  140. Platform: PlatformOpenAI,
  141. Window: "7d",
  142. ThresholdPercent: 80,
  143. UsedPercent: 91.5,
  144. Until: until,
  145. Now: until.Add(-time.Hour),
  146. })
  147. account := &Account{
  148. ID: 1002,
  149. Platform: PlatformOpenAI,
  150. Status: StatusActive,
  151. Schedulable: true,
  152. TempUnschedulableUntil: &until,
  153. TempUnschedulableReason: existingReason,
  154. Extra: map[string]any{
  155. "codex_7d_used_percent": 91.5,
  156. "codex_7d_reset_at": until.Format(time.RFC3339),
  157. },
  158. }
  159. blocked := rl.ApplyAccountSchedulingThreshold(context.Background(), account)
  160. require.True(t, blocked)
  161. require.Equal(t, 0, accountRepo.tempCalls)
  162. require.Equal(t, existingReason, account.TempUnschedulableReason)
  163. require.NotNil(t, account.TempUnschedulableUntil)
  164. require.True(t, until.Equal(*account.TempUnschedulableUntil))
  165. }
  166. func TestRateLimitService_ApplyAccountSchedulingThreshold_UnsupportedPlatformDoesNotBlock(t *testing.T) {
  167. accountSchedulingThresholdsSF.Forget(SettingKeyAccountSchedulingThresholds)
  168. accountSchedulingThresholdsCache.Store(&cachedAccountSchedulingThresholds{})
  169. settingsRepo := newMockSettingRepo()
  170. settingsRepo.data[SettingKeyAccountSchedulingThresholds] = `{"openai":80}`
  171. accountRepo := &rateLimitAccountRepoStub{}
  172. rl := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
  173. rl.SetSettingService(NewSettingService(settingsRepo, &config.Config{}))
  174. account := &Account{
  175. ID: 2002,
  176. Platform: PlatformKiro,
  177. Status: StatusActive,
  178. Schedulable: true,
  179. Credentials: map[string]any{
  180. "account_scheduling_threshold": 1,
  181. },
  182. Extra: map[string]any{
  183. "kiro_sched_utilization": 99.0,
  184. "kiro_sched_reset_at": time.Now().UTC().Add(24 * time.Hour).Format(time.RFC3339),
  185. },
  186. }
  187. blocked := rl.ApplyAccountSchedulingThreshold(context.Background(), account)
  188. require.False(t, blocked)
  189. require.Equal(t, 0, accountRepo.tempCalls)
  190. require.Nil(t, account.TempUnschedulableUntil)
  191. require.Empty(t, account.TempUnschedulableReason)
  192. }