|
- package channel
-
- import (
- "net/http"
- "net/http/httptest"
- "testing"
-
- relaycommon "github.com/QuantumNous/new-api/relay/common"
- "github.com/QuantumNous/new-api/setting/model_setting"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- )
-
- func TestApplyPassthroughModelMapping(t *testing.T) {
- body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
- info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: true, UpstreamModelName: "gpt-4.1"}}
-
- origRequestEnabled := model_setting.GetGlobalSettings().PassThroughRequestEnabled
- origApplyMapping := model_setting.GetGlobalSettings().PassthroughApplyModelMapping
- defer func() {
- model_setting.GetGlobalSettings().PassThroughRequestEnabled = origRequestEnabled
- model_setting.GetGlobalSettings().PassthroughApplyModelMapping = origApplyMapping
- }()
-
- model_setting.GetGlobalSettings().PassThroughRequestEnabled = true
- model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true
-
- got := applyPassthroughModelMapping(body, info)
- assert.JSONEq(t, `{"model":"gpt-4.1","messages":[{"role":"user","content":"hi"}]}`, string(got))
- }
-
- func TestApplyPassthroughModelMapping_NonJSONKeepsOriginal(t *testing.T) {
- body := []byte(`model=gpt-4o`)
- info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: true, UpstreamModelName: "gpt-4.1"}}
-
- origRequestEnabled := model_setting.GetGlobalSettings().PassThroughRequestEnabled
- origApplyMapping := model_setting.GetGlobalSettings().PassthroughApplyModelMapping
- defer func() {
- model_setting.GetGlobalSettings().PassThroughRequestEnabled = origRequestEnabled
- model_setting.GetGlobalSettings().PassthroughApplyModelMapping = origApplyMapping
- }()
-
- model_setting.GetGlobalSettings().PassThroughRequestEnabled = true
- model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true
-
- got := applyPassthroughModelMapping(body, info)
- assert.Equal(t, body, got)
- }
-
- func TestApplyPassthroughModelMapping_DisabledKeepsOriginal(t *testing.T) {
- body := []byte(`{"model":"gpt-4o"}`)
- info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: false, UpstreamModelName: "gpt-4.1"}}
-
- origRequestEnabled := model_setting.GetGlobalSettings().PassThroughRequestEnabled
- origApplyMapping := model_setting.GetGlobalSettings().PassthroughApplyModelMapping
- defer func() {
- model_setting.GetGlobalSettings().PassThroughRequestEnabled = origRequestEnabled
- model_setting.GetGlobalSettings().PassthroughApplyModelMapping = origApplyMapping
- }()
-
- model_setting.GetGlobalSettings().PassThroughRequestEnabled = true
- model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true
-
- got := applyPassthroughModelMapping(body, info)
- assert.Equal(t, body, got)
- }
-
- func TestDoApiRequest_DoesNotChangeNonPassthroughBody(t *testing.T) {
- origRequestEnabled := model_setting.GetGlobalSettings().PassThroughRequestEnabled
- origApplyMapping := model_setting.GetGlobalSettings().PassthroughApplyModelMapping
- defer func() {
- model_setting.GetGlobalSettings().PassThroughRequestEnabled = origRequestEnabled
- model_setting.GetGlobalSettings().PassthroughApplyModelMapping = origApplyMapping
- }()
-
- model_setting.GetGlobalSettings().PassThroughRequestEnabled = false
- model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true
-
- body := []byte(`{"model":"gpt-4o"}`)
- got := applyPassthroughModelMapping(body, &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: true, UpstreamModelName: "gpt-4.1"}})
- assert.Equal(t, body, got)
- }
-
- func TestProcessHeaderOverride_ChannelTestSkipsPassthroughRules(t *testing.T) {
- t.Parallel()
-
- gin.SetMode(gin.TestMode)
- recorder := httptest.NewRecorder()
- ctx, _ := gin.CreateTestContext(recorder)
- ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
- ctx.Request.Header.Set("X-Trace-Id", "trace-123")
-
- info := &relaycommon.RelayInfo{
- IsChannelTest: true,
- ChannelMeta: &relaycommon.ChannelMeta{
- HeadersOverride: map[string]any{
- "*": "",
- },
- },
- }
-
- headers, err := processHeaderOverride(info, ctx)
- require.NoError(t, err)
- require.Empty(t, headers)
- }
-
- func TestProcessHeaderOverride_ChannelTestSkipsClientHeaderPlaceholder(t *testing.T) {
- t.Parallel()
-
- gin.SetMode(gin.TestMode)
- recorder := httptest.NewRecorder()
- ctx, _ := gin.CreateTestContext(recorder)
- ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
- ctx.Request.Header.Set("X-Trace-Id", "trace-123")
-
- info := &relaycommon.RelayInfo{
- IsChannelTest: true,
- ChannelMeta: &relaycommon.ChannelMeta{
- HeadersOverride: map[string]any{
- "X-Upstream-Trace": "{client_header:X-Trace-Id}",
- },
- },
- }
-
- headers, err := processHeaderOverride(info, ctx)
- require.NoError(t, err)
- _, ok := headers["X-Upstream-Trace"]
- require.False(t, ok)
- }
-
- func TestProcessHeaderOverride_NonTestKeepsClientHeaderPlaceholder(t *testing.T) {
- t.Parallel()
-
- gin.SetMode(gin.TestMode)
- recorder := httptest.NewRecorder()
- ctx, _ := gin.CreateTestContext(recorder)
- ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
- ctx.Request.Header.Set("X-Trace-Id", "trace-123")
-
- info := &relaycommon.RelayInfo{
- IsChannelTest: false,
- ChannelMeta: &relaycommon.ChannelMeta{
- HeadersOverride: map[string]any{
- "X-Upstream-Trace": "{client_header:X-Trace-Id}",
- },
- },
- }
-
- headers, err := processHeaderOverride(info, ctx)
- require.NoError(t, err)
- require.Equal(t, "trace-123", headers["X-Upstream-Trace"])
- }
-
- func TestProcessHeaderOverride_PassthroughSkipsAcceptEncoding(t *testing.T) {
- t.Parallel()
-
- gin.SetMode(gin.TestMode)
- recorder := httptest.NewRecorder()
- ctx, _ := gin.CreateTestContext(recorder)
- ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
- ctx.Request.Header.Set("X-Trace-Id", "trace-123")
- ctx.Request.Header.Set("Accept-Encoding", "gzip")
-
- info := &relaycommon.RelayInfo{
- IsChannelTest: false,
- ChannelMeta: &relaycommon.ChannelMeta{
- HeadersOverride: map[string]any{
- "*": "",
- },
- },
- }
-
- headers, err := processHeaderOverride(info, ctx)
- require.NoError(t, err)
- require.Equal(t, "trace-123", headers["X-Trace-Id"])
-
- _, hasAcceptEncoding := headers["Accept-Encoding"]
- require.False(t, hasAcceptEncoding)
- }
|