|
|
|
@@ -0,0 +1,151 @@ |
|
|
|
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() *gin.Context { |
|
|
|
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() |
|
|
|
assert.False(t, shouldRetry(c, nil, 3)) |
|
|
|
} |
|
|
|
|
|
|
|
func TestShouldRetry_5xxErrorWithRetryLeft_ReturnsTrue(t *testing.T) { |
|
|
|
// 渠道选择删除后,所有请求都能正常重试 |
|
|
|
// 之前设置了 specific_channel_id 时 shouldRetry 会返回 false |
|
|
|
// 现在即使不设置任何 channel 相关 context,5xx 也应该重试 |
|
|
|
c := newTestContext() |
|
|
|
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() |
|
|
|
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() |
|
|
|
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() |
|
|
|
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() |
|
|
|
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() |
|
|
|
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() |
|
|
|
assert.False(t, shouldRetryTaskRelay(c, 1, nil, 3)) |
|
|
|
} |
|
|
|
|
|
|
|
func TestShouldRetryTaskRelay_5xxWithRetryLeft_ReturnsTrue(t *testing.T) { |
|
|
|
c := newTestContext() |
|
|
|
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() |
|
|
|
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() |
|
|
|
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() |
|
|
|
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() |
|
|
|
taskErr := &dto.TaskError{ |
|
|
|
Error: errors.New("local error"), |
|
|
|
StatusCode: http.StatusForbidden, |
|
|
|
LocalError: true, |
|
|
|
} |
|
|
|
assert.False(t, shouldRetryTaskRelay(c, 1, taskErr, 3)) |
|
|
|
} |