|
- 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))
- }
|