ratelimit_service_anthropic_test.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339
  1. package service
  2. import (
  3. "net/http"
  4. "strconv"
  5. "testing"
  6. "time"
  7. )
  8. func TestCalculateAnthropic429ResetTime_Only5hExceeded(t *testing.T) {
  9. headers := http.Header{}
  10. headers.Set("anthropic-ratelimit-unified-5h-utilization", "1.02")
  11. headers.Set("anthropic-ratelimit-unified-5h-reset", "1770998400")
  12. headers.Set("anthropic-ratelimit-unified-7d-utilization", "0.32")
  13. headers.Set("anthropic-ratelimit-unified-7d-reset", "1771549200")
  14. result := calculateAnthropic429ResetTime(headers)
  15. assertAnthropicResult(t, result, 1770998400)
  16. if result.fiveHourReset == nil || !result.fiveHourReset.Equal(time.Unix(1770998400, 0)) {
  17. t.Errorf("expected fiveHourReset=1770998400, got %v", result.fiveHourReset)
  18. }
  19. }
  20. func TestCalculateAnthropic429ResetTime_Only7dExceeded(t *testing.T) {
  21. headers := http.Header{}
  22. headers.Set("anthropic-ratelimit-unified-5h-utilization", "0.50")
  23. headers.Set("anthropic-ratelimit-unified-5h-reset", "1770998400")
  24. headers.Set("anthropic-ratelimit-unified-7d-utilization", "1.05")
  25. headers.Set("anthropic-ratelimit-unified-7d-reset", "1771549200")
  26. result := calculateAnthropic429ResetTime(headers)
  27. assertAnthropicResult(t, result, 1771549200)
  28. // fiveHourReset should still be populated for session window calculation
  29. if result.fiveHourReset == nil || !result.fiveHourReset.Equal(time.Unix(1770998400, 0)) {
  30. t.Errorf("expected fiveHourReset=1770998400, got %v", result.fiveHourReset)
  31. }
  32. }
  33. func TestCalculateAnthropic429ResetTime_BothExceeded(t *testing.T) {
  34. headers := http.Header{}
  35. headers.Set("anthropic-ratelimit-unified-5h-utilization", "1.10")
  36. headers.Set("anthropic-ratelimit-unified-5h-reset", "1770998400")
  37. headers.Set("anthropic-ratelimit-unified-7d-utilization", "1.02")
  38. headers.Set("anthropic-ratelimit-unified-7d-reset", "1771549200")
  39. result := calculateAnthropic429ResetTime(headers)
  40. assertAnthropicResult(t, result, 1771549200)
  41. }
  42. func TestCalculateAnthropic429ResetTime_NoPerWindowHeaders(t *testing.T) {
  43. headers := http.Header{}
  44. headers.Set("anthropic-ratelimit-unified-reset", "1771549200")
  45. result := calculateAnthropic429ResetTime(headers)
  46. if result != nil {
  47. t.Errorf("expected nil result when no per-window headers, got resetAt=%v", result.resetAt)
  48. }
  49. }
  50. func TestCalculateAnthropic429ResetTime_NoHeaders(t *testing.T) {
  51. result := calculateAnthropic429ResetTime(http.Header{})
  52. if result != nil {
  53. t.Errorf("expected nil result for empty headers, got resetAt=%v", result.resetAt)
  54. }
  55. }
  56. func TestCalculateAnthropic429ResetTime_SurpassedThreshold(t *testing.T) {
  57. headers := http.Header{}
  58. headers.Set("anthropic-ratelimit-unified-5h-surpassed-threshold", "true")
  59. headers.Set("anthropic-ratelimit-unified-5h-reset", "1770998400")
  60. headers.Set("anthropic-ratelimit-unified-7d-surpassed-threshold", "false")
  61. headers.Set("anthropic-ratelimit-unified-7d-reset", "1771549200")
  62. result := calculateAnthropic429ResetTime(headers)
  63. assertAnthropicResult(t, result, 1770998400)
  64. }
  65. func TestCalculateAnthropic429ResetTime_UtilizationExactlyOne(t *testing.T) {
  66. headers := http.Header{}
  67. headers.Set("anthropic-ratelimit-unified-5h-utilization", "1.0")
  68. headers.Set("anthropic-ratelimit-unified-5h-reset", "1770998400")
  69. headers.Set("anthropic-ratelimit-unified-7d-utilization", "0.5")
  70. headers.Set("anthropic-ratelimit-unified-7d-reset", "1771549200")
  71. result := calculateAnthropic429ResetTime(headers)
  72. assertAnthropicResult(t, result, 1770998400)
  73. }
  74. func TestCalculateAnthropic429ResetTime_NeitherExceeded_UsesShorter(t *testing.T) {
  75. headers := http.Header{}
  76. headers.Set("anthropic-ratelimit-unified-5h-utilization", "0.95")
  77. headers.Set("anthropic-ratelimit-unified-5h-reset", "1770998400") // sooner
  78. headers.Set("anthropic-ratelimit-unified-7d-utilization", "0.80")
  79. headers.Set("anthropic-ratelimit-unified-7d-reset", "1771549200") // later
  80. result := calculateAnthropic429ResetTime(headers)
  81. assertAnthropicResult(t, result, 1770998400)
  82. }
  83. func TestCalculateAnthropic429ResetTime_Only5hResetHeader(t *testing.T) {
  84. headers := http.Header{}
  85. headers.Set("anthropic-ratelimit-unified-5h-utilization", "1.05")
  86. headers.Set("anthropic-ratelimit-unified-5h-reset", "1770998400")
  87. result := calculateAnthropic429ResetTime(headers)
  88. assertAnthropicResult(t, result, 1770998400)
  89. }
  90. func TestCalculateAnthropic429ResetTime_Only7dResetHeader(t *testing.T) {
  91. headers := http.Header{}
  92. headers.Set("anthropic-ratelimit-unified-7d-utilization", "1.03")
  93. headers.Set("anthropic-ratelimit-unified-7d-reset", "1771549200")
  94. result := calculateAnthropic429ResetTime(headers)
  95. assertAnthropicResult(t, result, 1771549200)
  96. if result.fiveHourReset != nil {
  97. t.Errorf("expected fiveHourReset=nil when no 5h headers, got %v", result.fiveHourReset)
  98. }
  99. }
  100. func TestIsAnthropicWindowExceeded(t *testing.T) {
  101. tests := []struct {
  102. name string
  103. headers http.Header
  104. window string
  105. expected bool
  106. }{
  107. {
  108. name: "utilization above 1.0",
  109. headers: makeHeader("anthropic-ratelimit-unified-5h-utilization", "1.02"),
  110. window: "5h",
  111. expected: true,
  112. },
  113. {
  114. name: "utilization exactly 1.0",
  115. headers: makeHeader("anthropic-ratelimit-unified-5h-utilization", "1.0"),
  116. window: "5h",
  117. expected: true,
  118. },
  119. {
  120. name: "utilization below 1.0",
  121. headers: makeHeader("anthropic-ratelimit-unified-5h-utilization", "0.99"),
  122. window: "5h",
  123. expected: false,
  124. },
  125. {
  126. name: "surpassed-threshold true",
  127. headers: makeHeader("anthropic-ratelimit-unified-7d-surpassed-threshold", "true"),
  128. window: "7d",
  129. expected: true,
  130. },
  131. {
  132. name: "surpassed-threshold True (case insensitive)",
  133. headers: makeHeader("anthropic-ratelimit-unified-7d-surpassed-threshold", "True"),
  134. window: "7d",
  135. expected: true,
  136. },
  137. {
  138. name: "surpassed-threshold false",
  139. headers: makeHeader("anthropic-ratelimit-unified-7d-surpassed-threshold", "false"),
  140. window: "7d",
  141. expected: false,
  142. },
  143. {
  144. name: "no headers",
  145. headers: http.Header{},
  146. window: "5h",
  147. expected: false,
  148. },
  149. }
  150. for _, tc := range tests {
  151. t.Run(tc.name, func(t *testing.T) {
  152. got := isAnthropicWindowExceeded(tc.headers, tc.window)
  153. if got != tc.expected {
  154. t.Errorf("expected %v, got %v", tc.expected, got)
  155. }
  156. })
  157. }
  158. }
  159. func TestSelectAnthropicFableWindowLimit_RejectedStatus(t *testing.T) {
  160. now := time.Now()
  161. reset := now.Add(80 * time.Hour).Truncate(time.Second)
  162. headers := http.Header{}
  163. headers.Set("anthropic-ratelimit-unified-7d_oi-status", "rejected")
  164. headers.Set("anthropic-ratelimit-unified-7d_oi-utilization", "1.0")
  165. headers.Set("anthropic-ratelimit-unified-7d_oi-surpassed-threshold", "1.0")
  166. headers.Set("anthropic-ratelimit-unified-7d_oi-reset", strconv.FormatInt(reset.Unix(), 10))
  167. limit := selectAnthropicFableWindowLimit(headers, now)
  168. if limit == nil {
  169. t.Fatal("expected non-nil limit")
  170. }
  171. if !limit.resetAt.Equal(reset) {
  172. t.Errorf("expected resetAt=%v, got %v", reset, limit.resetAt)
  173. }
  174. if limit.reason != anthropicFableWindowReason {
  175. t.Errorf("expected reason=%q, got %q", anthropicFableWindowReason, limit.reason)
  176. }
  177. }
  178. func TestSelectAnthropicFableWindowLimit_UtilizationOnly(t *testing.T) {
  179. // 无 status 头时,utilization >= 1.0 也应视为超限
  180. now := time.Now()
  181. reset := now.Add(3 * 24 * time.Hour).Truncate(time.Second)
  182. headers := http.Header{}
  183. headers.Set("anthropic-ratelimit-unified-7d_oi-utilization", "1.0")
  184. headers.Set("anthropic-ratelimit-unified-7d_oi-reset", strconv.FormatInt(reset.Unix(), 10))
  185. limit := selectAnthropicFableWindowLimit(headers, now)
  186. if limit == nil {
  187. t.Fatal("expected non-nil limit")
  188. }
  189. if !limit.resetAt.Equal(reset) {
  190. t.Errorf("expected resetAt=%v, got %v", reset, limit.resetAt)
  191. }
  192. }
  193. func TestSelectAnthropicFableWindowLimit_AllowedReturnsNil(t *testing.T) {
  194. now := time.Now()
  195. headers := http.Header{}
  196. headers.Set("anthropic-ratelimit-unified-7d_oi-status", "allowed")
  197. headers.Set("anthropic-ratelimit-unified-7d_oi-utilization", "0.56")
  198. headers.Set("anthropic-ratelimit-unified-7d_oi-reset", strconv.FormatInt(now.Add(80*time.Hour).Unix(), 10))
  199. if limit := selectAnthropicFableWindowLimit(headers, now); limit != nil {
  200. t.Errorf("expected nil limit for allowed window, got %+v", limit)
  201. }
  202. }
  203. func TestSelectAnthropicFableWindowLimit_NoHeadersReturnsNil(t *testing.T) {
  204. if limit := selectAnthropicFableWindowLimit(http.Header{}, time.Now()); limit != nil {
  205. t.Errorf("expected nil limit for empty headers, got %+v", limit)
  206. }
  207. }
  208. func TestSelectAnthropicFableWindowLimit_FallsBackToAggregateReset(t *testing.T) {
  209. // 7d_oi-reset 缺失时回退聚合 anthropic-ratelimit-unified-reset
  210. now := time.Now()
  211. reset := now.Add(80 * time.Hour).Truncate(time.Second)
  212. headers := http.Header{}
  213. headers.Set("anthropic-ratelimit-unified-7d_oi-status", "rejected")
  214. headers.Set("anthropic-ratelimit-unified-reset", strconv.FormatInt(reset.Unix(), 10))
  215. limit := selectAnthropicFableWindowLimit(headers, now)
  216. if limit == nil {
  217. t.Fatal("expected non-nil limit via aggregate reset fallback")
  218. }
  219. if !limit.resetAt.Equal(reset) {
  220. t.Errorf("expected resetAt=%v, got %v", reset, limit.resetAt)
  221. }
  222. }
  223. func TestSelectAnthropicFableWindowLimit_RejectedWithoutAnyResetReturnsNil(t *testing.T) {
  224. headers := http.Header{}
  225. headers.Set("anthropic-ratelimit-unified-7d_oi-status", "rejected")
  226. if limit := selectAnthropicFableWindowLimit(headers, time.Now()); limit != nil {
  227. t.Errorf("expected nil limit when no reset time available, got %+v", limit)
  228. }
  229. }
  230. func TestParseAnthropicAggregateReset(t *testing.T) {
  231. now := time.Now()
  232. future := now.Add(80 * time.Hour).Truncate(time.Second)
  233. tests := []struct {
  234. name string
  235. value string
  236. want time.Time
  237. wantOK bool
  238. }{
  239. {"valid seconds", strconv.FormatInt(future.Unix(), 10), future, true},
  240. {"valid milliseconds", strconv.FormatInt(future.UnixMilli(), 10), future, true},
  241. {"empty", "", time.Time{}, false},
  242. {"garbage", "abc", time.Time{}, false},
  243. {"in the past", strconv.FormatInt(now.Add(-time.Hour).Unix(), 10), time.Time{}, false},
  244. {"too far in the future", strconv.FormatInt(now.Add(30*24*time.Hour).Unix(), 10), time.Time{}, false},
  245. }
  246. for _, tc := range tests {
  247. t.Run(tc.name, func(t *testing.T) {
  248. headers := http.Header{}
  249. if tc.value != "" {
  250. headers.Set("anthropic-ratelimit-unified-reset", tc.value)
  251. }
  252. got, ok := parseAnthropicAggregateReset(headers, now)
  253. if ok != tc.wantOK {
  254. t.Fatalf("expected ok=%v, got %v", tc.wantOK, ok)
  255. }
  256. if ok && !got.Equal(tc.want) {
  257. t.Errorf("expected %v, got %v", tc.want, got)
  258. }
  259. })
  260. }
  261. }
  262. func TestIsAnthropicWindowRejected(t *testing.T) {
  263. headers := http.Header{}
  264. headers.Set("anthropic-ratelimit-unified-7d_oi-status", "Rejected")
  265. headers.Set("anthropic-ratelimit-unified-5h-status", "allowed")
  266. if !isAnthropicWindowRejected(headers, "7d_oi") {
  267. t.Error("expected 7d_oi to be rejected (case insensitive)")
  268. }
  269. if isAnthropicWindowRejected(headers, "5h") {
  270. t.Error("expected 5h not rejected")
  271. }
  272. if isAnthropicWindowRejected(headers, "7d") {
  273. t.Error("expected missing 7d status not rejected")
  274. }
  275. }
  276. // assertAnthropicResult is a test helper that verifies the result is non-nil and
  277. // has the expected resetAt unix timestamp.
  278. func assertAnthropicResult(t *testing.T, result *anthropic429Result, wantUnix int64) {
  279. t.Helper()
  280. if result == nil {
  281. t.Fatal("expected non-nil result")
  282. return // unreachable, but satisfies staticcheck SA5011
  283. }
  284. want := time.Unix(wantUnix, 0)
  285. if !result.resetAt.Equal(want) {
  286. t.Errorf("expected resetAt=%v, got %v", want, result.resetAt)
  287. }
  288. }
  289. func makeHeader(key, value string) http.Header {
  290. h := http.Header{}
  291. h.Set(key, value)
  292. return h
  293. }