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

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