package controller import ( "errors" "net/http" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" ) // --------------------------------------------------------------------------- // shouldRetry — 渠道选择删除后的验证测试 // --------------------------------------------------------------------------- func newTestContext(t *testing.T) *gin.Context { t.Helper() gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request, _ = http.NewRequest("POST", "/v1/chat/completions", nil) return c } func TestShouldRetry_NilError_ReturnsFalse(t *testing.T) { c := newTestContext(t) assert.False(t, shouldRetry(c, nil, 3)) } func TestShouldRetry_5xxErrorWithRetryLeft_ReturnsTrue(t *testing.T) { // 渠道选择删除后,所有请求都能正常重试 // 之前设置了 specific_channel_id 时 shouldRetry 会返回 false // 现在即使不设置任何 channel 相关 context,5xx 也应该重试 c := newTestContext(t) apiErr := types.NewErrorWithStatusCode( errors.New("upstream error"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError, ) assert.True(t, shouldRetry(c, apiErr, 3)) } func TestShouldRetry_5xxErrorZeroRetryLeft_ReturnsFalse(t *testing.T) { c := newTestContext(t) apiErr := types.NewErrorWithStatusCode( errors.New("upstream error"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError, ) assert.False(t, shouldRetry(c, apiErr, 0)) } func TestShouldRetry_SkipRetryError_ReturnsFalse(t *testing.T) { c := newTestContext(t) apiErr := types.NewError( errors.New("skip retry"), types.ErrorCodeBadResponseStatusCode, types.ErrOptionWithSkipRetry(), ) assert.False(t, shouldRetry(c, apiErr, 3)) } func TestShouldRetry_ChannelError_ReturnsTrue(t *testing.T) { // channel: 前缀的错误码会触发立即重试 c := newTestContext(t) apiErr := types.NewError( errors.New("channel error"), "channel:test", ) // 即使 retryTimes=0,channel error 也返回 true assert.True(t, shouldRetry(c, apiErr, 0)) } func TestShouldRetry_2xxStatus_ReturnsFalse(t *testing.T) { c := newTestContext(t) apiErr := types.NewErrorWithStatusCode( errors.New("ok"), types.ErrorCodeBadResponseStatusCode, http.StatusOK, ) assert.False(t, shouldRetry(c, apiErr, 3)) } func TestShouldRetry_429Status_ReturnsTrue(t *testing.T) { c := newTestContext(t) apiErr := types.NewErrorWithStatusCode( errors.New("rate limited"), types.ErrorCodeBadResponseStatusCode, http.StatusTooManyRequests, ) assert.True(t, shouldRetry(c, apiErr, 3)) } // --------------------------------------------------------------------------- // shouldRetryTaskRelay — 渠道选择删除后的验证测试 // --------------------------------------------------------------------------- func TestShouldRetryTaskRelay_NilError_ReturnsFalse(t *testing.T) { c := newTestContext(t) assert.False(t, shouldRetryTaskRelay(c, 1, nil, 3)) } func TestShouldRetryTaskRelay_5xxWithRetryLeft_ReturnsTrue(t *testing.T) { c := newTestContext(t) taskErr := &dto.TaskError{ Error: errors.New("server error"), StatusCode: http.StatusInternalServerError, } assert.True(t, shouldRetryTaskRelay(c, 1, taskErr, 3)) } func TestShouldRetryTaskRelay_5xxZeroRetryLeft_ReturnsFalse(t *testing.T) { c := newTestContext(t) taskErr := &dto.TaskError{ Error: errors.New("server error"), StatusCode: http.StatusInternalServerError, } assert.False(t, shouldRetryTaskRelay(c, 1, taskErr, 0)) } func TestShouldRetryTaskRelay_429_ReturnsTrue(t *testing.T) { c := newTestContext(t) taskErr := &dto.TaskError{ Error: errors.New("rate limited"), StatusCode: http.StatusTooManyRequests, } assert.True(t, shouldRetryTaskRelay(c, 1, taskErr, 1)) } func TestShouldRetryTaskRelay_400_ReturnsFalse(t *testing.T) { c := newTestContext(t) taskErr := &dto.TaskError{ Error: errors.New("bad request"), StatusCode: http.StatusBadRequest, } assert.False(t, shouldRetryTaskRelay(c, 1, taskErr, 3)) } func TestShouldRetryTaskRelay_LocalError_ReturnsFalse(t *testing.T) { // LocalError + 非 5xx 状态码(如 403)时,LocalError 检查生效返回 false // 注意:5xx 分支在 LocalError 检查之前,所以 localError+5xx 仍会重试 c := newTestContext(t) taskErr := &dto.TaskError{ Error: errors.New("local error"), StatusCode: http.StatusForbidden, LocalError: true, } assert.False(t, shouldRetryTaskRelay(c, 1, taskErr, 3)) }