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 TestPassthroughModelMapping_AppliedWhenEnabled(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 = true model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: true, UpstreamModelName: "gpt-4.1"}} body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`) assert.True(t, needsPassthroughModelMapping(info)) got := applyPassthroughModelMapping(body, info) assert.JSONEq(t, `{"model":"gpt-4.1","messages":[{"role":"user","content":"hi"}]}`, string(got)) } func TestPassthroughModelMapping_SkippedWhenNotMapped(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 = true model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: false, UpstreamModelName: "gpt-4.1"}} assert.False(t, needsPassthroughModelMapping(info)) } func TestPassthroughModelMapping_SkippedWhenPassthroughOff(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 info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: true, UpstreamModelName: "gpt-4.1"}} assert.False(t, needsPassthroughModelMapping(info)) } func TestPassthroughModelMapping_NonJSONKeepsOriginal(t *testing.T) { body := []byte(`model=gpt-4o`) info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: true, UpstreamModelName: "gpt-4.1"}} got := applyPassthroughModelMapping(body, info) 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) }