Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

153 строки
4.5 KiB

  1. package controller
  2. import (
  3. "errors"
  4. "net/http"
  5. "net/http/httptest"
  6. "testing"
  7. "github.com/QuantumNous/new-api/dto"
  8. "github.com/QuantumNous/new-api/types"
  9. "github.com/gin-gonic/gin"
  10. "github.com/stretchr/testify/assert"
  11. )
  12. // ---------------------------------------------------------------------------
  13. // shouldRetry — 渠道选择删除后的验证测试
  14. // ---------------------------------------------------------------------------
  15. func newTestContext(t *testing.T) *gin.Context {
  16. t.Helper()
  17. gin.SetMode(gin.TestMode)
  18. w := httptest.NewRecorder()
  19. c, _ := gin.CreateTestContext(w)
  20. c.Request, _ = http.NewRequest("POST", "/v1/chat/completions", nil)
  21. return c
  22. }
  23. func TestShouldRetry_NilError_ReturnsFalse(t *testing.T) {
  24. c := newTestContext(t)
  25. assert.False(t, shouldRetry(c, nil, 3))
  26. }
  27. func TestShouldRetry_5xxErrorWithRetryLeft_ReturnsTrue(t *testing.T) {
  28. // 渠道选择删除后,所有请求都能正常重试
  29. // 之前设置了 specific_channel_id 时 shouldRetry 会返回 false
  30. // 现在即使不设置任何 channel 相关 context,5xx 也应该重试
  31. c := newTestContext(t)
  32. apiErr := types.NewErrorWithStatusCode(
  33. errors.New("upstream error"),
  34. types.ErrorCodeBadResponseStatusCode,
  35. http.StatusInternalServerError,
  36. )
  37. assert.True(t, shouldRetry(c, apiErr, 3))
  38. }
  39. func TestShouldRetry_5xxErrorZeroRetryLeft_ReturnsFalse(t *testing.T) {
  40. c := newTestContext(t)
  41. apiErr := types.NewErrorWithStatusCode(
  42. errors.New("upstream error"),
  43. types.ErrorCodeBadResponseStatusCode,
  44. http.StatusInternalServerError,
  45. )
  46. assert.False(t, shouldRetry(c, apiErr, 0))
  47. }
  48. func TestShouldRetry_SkipRetryError_ReturnsFalse(t *testing.T) {
  49. c := newTestContext(t)
  50. apiErr := types.NewError(
  51. errors.New("skip retry"),
  52. types.ErrorCodeBadResponseStatusCode,
  53. types.ErrOptionWithSkipRetry(),
  54. )
  55. assert.False(t, shouldRetry(c, apiErr, 3))
  56. }
  57. func TestShouldRetry_ChannelError_ReturnsTrue(t *testing.T) {
  58. // channel: 前缀的错误码会触发立即重试
  59. c := newTestContext(t)
  60. apiErr := types.NewError(
  61. errors.New("channel error"),
  62. "channel:test",
  63. )
  64. // 即使 retryTimes=0,channel error 也返回 true
  65. assert.True(t, shouldRetry(c, apiErr, 0))
  66. }
  67. func TestShouldRetry_2xxStatus_ReturnsFalse(t *testing.T) {
  68. c := newTestContext(t)
  69. apiErr := types.NewErrorWithStatusCode(
  70. errors.New("ok"),
  71. types.ErrorCodeBadResponseStatusCode,
  72. http.StatusOK,
  73. )
  74. assert.False(t, shouldRetry(c, apiErr, 3))
  75. }
  76. func TestShouldRetry_429Status_ReturnsTrue(t *testing.T) {
  77. c := newTestContext(t)
  78. apiErr := types.NewErrorWithStatusCode(
  79. errors.New("rate limited"),
  80. types.ErrorCodeBadResponseStatusCode,
  81. http.StatusTooManyRequests,
  82. )
  83. assert.True(t, shouldRetry(c, apiErr, 3))
  84. }
  85. // ---------------------------------------------------------------------------
  86. // shouldRetryTaskRelay — 渠道选择删除后的验证测试
  87. // ---------------------------------------------------------------------------
  88. func TestShouldRetryTaskRelay_NilError_ReturnsFalse(t *testing.T) {
  89. c := newTestContext(t)
  90. assert.False(t, shouldRetryTaskRelay(c, 1, nil, 3))
  91. }
  92. func TestShouldRetryTaskRelay_5xxWithRetryLeft_ReturnsTrue(t *testing.T) {
  93. c := newTestContext(t)
  94. taskErr := &dto.TaskError{
  95. Error: errors.New("server error"),
  96. StatusCode: http.StatusInternalServerError,
  97. }
  98. assert.True(t, shouldRetryTaskRelay(c, 1, taskErr, 3))
  99. }
  100. func TestShouldRetryTaskRelay_5xxZeroRetryLeft_ReturnsFalse(t *testing.T) {
  101. c := newTestContext(t)
  102. taskErr := &dto.TaskError{
  103. Error: errors.New("server error"),
  104. StatusCode: http.StatusInternalServerError,
  105. }
  106. assert.False(t, shouldRetryTaskRelay(c, 1, taskErr, 0))
  107. }
  108. func TestShouldRetryTaskRelay_429_ReturnsTrue(t *testing.T) {
  109. c := newTestContext(t)
  110. taskErr := &dto.TaskError{
  111. Error: errors.New("rate limited"),
  112. StatusCode: http.StatusTooManyRequests,
  113. }
  114. assert.True(t, shouldRetryTaskRelay(c, 1, taskErr, 1))
  115. }
  116. func TestShouldRetryTaskRelay_400_ReturnsFalse(t *testing.T) {
  117. c := newTestContext(t)
  118. taskErr := &dto.TaskError{
  119. Error: errors.New("bad request"),
  120. StatusCode: http.StatusBadRequest,
  121. }
  122. assert.False(t, shouldRetryTaskRelay(c, 1, taskErr, 3))
  123. }
  124. func TestShouldRetryTaskRelay_LocalError_ReturnsFalse(t *testing.T) {
  125. // LocalError + 非 5xx 状态码(如 403)时,LocalError 检查生效返回 false
  126. // 注意:5xx 分支在 LocalError 检查之前,所以 localError+5xx 仍会重试
  127. c := newTestContext(t)
  128. taskErr := &dto.TaskError{
  129. Error: errors.New("local error"),
  130. StatusCode: http.StatusForbidden,
  131. LocalError: true,
  132. }
  133. assert.False(t, shouldRetryTaskRelay(c, 1, taskErr, 3))
  134. }