Просмотр исходного кода

test: 添加渠道选择功能删除后的验证测试

验证 parseTokenKey 不再解析 channelId(sk-abc123:42 中的 :42
保留为 key 一部分),shouldRetry/shouldRetryTaskRelay 不再检查
specific_channel_id,所有请求都能正常重试。
master
fengsilin 4 недель назад
Родитель
Сommit
631e030411
2 измененных файлов: 200 добавлений и 0 удалений
  1. +151
    -0
      controller/relay_test.go
  2. +49
    -0
      middleware/auth_test.go

+ 151
- 0
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))
}

+ 49
- 0
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)
}

Загрузка…
Отмена
Сохранить