From 631e030411c989420b987b518c7c214b3eca4abe Mon Sep 17 00:00:00 2001 From: fengsilin Date: Sat, 16 May 2026 01:24:52 +0800 Subject: [PATCH] =?UTF-8?q?test:=20=E6=B7=BB=E5=8A=A0=E6=B8=A0=E9=81=93?= =?UTF-8?q?=E9=80=89=E6=8B=A9=E5=8A=9F=E8=83=BD=E5=88=A0=E9=99=A4=E5=90=8E?= =?UTF-8?q?=E7=9A=84=E9=AA=8C=E8=AF=81=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 验证 parseTokenKey 不再解析 channelId(sk-abc123:42 中的 :42 保留为 key 一部分),shouldRetry/shouldRetryTaskRelay 不再检查 specific_channel_id,所有请求都能正常重试。 --- controller/relay_test.go | 151 +++++++++++++++++++++++++++++++++++++++ middleware/auth_test.go | 49 +++++++++++++ 2 files changed, 200 insertions(+) create mode 100644 controller/relay_test.go create mode 100644 middleware/auth_test.go diff --git a/controller/relay_test.go b/controller/relay_test.go new file mode 100644 index 0000000..8130f54 --- /dev/null +++ b/controller/relay_test.go @@ -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)) +} diff --git a/middleware/auth_test.go b/middleware/auth_test.go new file mode 100644 index 0000000..d4e67b0 --- /dev/null +++ b/middleware/auth_test.go @@ -0,0 +1,49 @@ +package middleware + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// --------------------------------------------------------------------------- +// parseTokenKey — 渠道选择删除后的验证测试 +// --------------------------------------------------------------------------- + +func TestParseTokenKey_SimpleKey(t *testing.T) { + key, parts := parseTokenKey("sk-abc123") + assert.Equal(t, "abc123", key) + assert.Equal(t, []string{"abc123"}, parts) +} + +func TestParseTokenKey_KeyWithDash(t *testing.T) { + // "sk-abc123-42" 中 "-" 用于分割,parts[0]="abc123" 就是 key + // parts[1]="42" 不再被解释为 channelId(渠道选择功能已删除) + key, parts := parseTokenKey("sk-abc123-42") + assert.Equal(t, "abc123", key) + assert.Equal(t, []string{"abc123", "42"}, parts) + // 验证:parts[1] 存在但不会被用于渠道选择 + // 渠道选择功能已删除,parseTokenKey 仅返回 parts,不解析 channelId +} + +func TestParseTokenKey_ColonFormat(t *testing.T) { + // 旧格式 "sk-abc123:42" 中 ":" 不再被解析为 channelId + // parseTokenKey 只按 "-" 分割,":" 会保留在 parts[0] 中 + key, parts := parseTokenKey("sk-abc123:42") + assert.Equal(t, "abc123:42", key) + assert.Equal(t, []string{"abc123:42"}, parts) + // 旧行为会返回 key="abc123" 并提取 channelId=42 + // 新行为:":" 被当作 key 的一部分,不做特殊处理 +} + +func TestParseTokenKey_EmptyKey(t *testing.T) { + key, parts := parseTokenKey("sk-") + assert.Equal(t, "", key) + assert.Equal(t, []string{""}, parts) +} + +func TestParseTokenKey_MultipleDashes(t *testing.T) { + key, parts := parseTokenKey("sk-abc-def-ghi") + assert.Equal(t, "abc", key) + assert.Equal(t, []string{"abc", "def", "ghi"}, parts) +}