You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

181 lines
6.2 KiB

  1. package channel
  2. import (
  3. "net/http"
  4. "net/http/httptest"
  5. "testing"
  6. relaycommon "github.com/QuantumNous/new-api/relay/common"
  7. "github.com/QuantumNous/new-api/setting/model_setting"
  8. "github.com/gin-gonic/gin"
  9. "github.com/stretchr/testify/assert"
  10. "github.com/stretchr/testify/require"
  11. )
  12. func TestApplyPassthroughModelMapping(t *testing.T) {
  13. body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
  14. info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: true, UpstreamModelName: "gpt-4.1"}}
  15. origRequestEnabled := model_setting.GetGlobalSettings().PassThroughRequestEnabled
  16. origApplyMapping := model_setting.GetGlobalSettings().PassthroughApplyModelMapping
  17. defer func() {
  18. model_setting.GetGlobalSettings().PassThroughRequestEnabled = origRequestEnabled
  19. model_setting.GetGlobalSettings().PassthroughApplyModelMapping = origApplyMapping
  20. }()
  21. model_setting.GetGlobalSettings().PassThroughRequestEnabled = true
  22. model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true
  23. got := applyPassthroughModelMapping(body, info)
  24. assert.JSONEq(t, `{"model":"gpt-4.1","messages":[{"role":"user","content":"hi"}]}`, string(got))
  25. }
  26. func TestApplyPassthroughModelMapping_NonJSONKeepsOriginal(t *testing.T) {
  27. body := []byte(`model=gpt-4o`)
  28. info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: true, UpstreamModelName: "gpt-4.1"}}
  29. origRequestEnabled := model_setting.GetGlobalSettings().PassThroughRequestEnabled
  30. origApplyMapping := model_setting.GetGlobalSettings().PassthroughApplyModelMapping
  31. defer func() {
  32. model_setting.GetGlobalSettings().PassThroughRequestEnabled = origRequestEnabled
  33. model_setting.GetGlobalSettings().PassthroughApplyModelMapping = origApplyMapping
  34. }()
  35. model_setting.GetGlobalSettings().PassThroughRequestEnabled = true
  36. model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true
  37. got := applyPassthroughModelMapping(body, info)
  38. assert.Equal(t, body, got)
  39. }
  40. func TestApplyPassthroughModelMapping_DisabledKeepsOriginal(t *testing.T) {
  41. body := []byte(`{"model":"gpt-4o"}`)
  42. info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: false, UpstreamModelName: "gpt-4.1"}}
  43. origRequestEnabled := model_setting.GetGlobalSettings().PassThroughRequestEnabled
  44. origApplyMapping := model_setting.GetGlobalSettings().PassthroughApplyModelMapping
  45. defer func() {
  46. model_setting.GetGlobalSettings().PassThroughRequestEnabled = origRequestEnabled
  47. model_setting.GetGlobalSettings().PassthroughApplyModelMapping = origApplyMapping
  48. }()
  49. model_setting.GetGlobalSettings().PassThroughRequestEnabled = true
  50. model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true
  51. got := applyPassthroughModelMapping(body, info)
  52. assert.Equal(t, body, got)
  53. }
  54. func TestDoApiRequest_DoesNotChangeNonPassthroughBody(t *testing.T) {
  55. origRequestEnabled := model_setting.GetGlobalSettings().PassThroughRequestEnabled
  56. origApplyMapping := model_setting.GetGlobalSettings().PassthroughApplyModelMapping
  57. defer func() {
  58. model_setting.GetGlobalSettings().PassThroughRequestEnabled = origRequestEnabled
  59. model_setting.GetGlobalSettings().PassthroughApplyModelMapping = origApplyMapping
  60. }()
  61. model_setting.GetGlobalSettings().PassThroughRequestEnabled = false
  62. model_setting.GetGlobalSettings().PassthroughApplyModelMapping = true
  63. body := []byte(`{"model":"gpt-4o"}`)
  64. got := applyPassthroughModelMapping(body, &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{IsModelMapped: true, UpstreamModelName: "gpt-4.1"}})
  65. assert.Equal(t, body, got)
  66. }
  67. func TestProcessHeaderOverride_ChannelTestSkipsPassthroughRules(t *testing.T) {
  68. t.Parallel()
  69. gin.SetMode(gin.TestMode)
  70. recorder := httptest.NewRecorder()
  71. ctx, _ := gin.CreateTestContext(recorder)
  72. ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
  73. ctx.Request.Header.Set("X-Trace-Id", "trace-123")
  74. info := &relaycommon.RelayInfo{
  75. IsChannelTest: true,
  76. ChannelMeta: &relaycommon.ChannelMeta{
  77. HeadersOverride: map[string]any{
  78. "*": "",
  79. },
  80. },
  81. }
  82. headers, err := processHeaderOverride(info, ctx)
  83. require.NoError(t, err)
  84. require.Empty(t, headers)
  85. }
  86. func TestProcessHeaderOverride_ChannelTestSkipsClientHeaderPlaceholder(t *testing.T) {
  87. t.Parallel()
  88. gin.SetMode(gin.TestMode)
  89. recorder := httptest.NewRecorder()
  90. ctx, _ := gin.CreateTestContext(recorder)
  91. ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
  92. ctx.Request.Header.Set("X-Trace-Id", "trace-123")
  93. info := &relaycommon.RelayInfo{
  94. IsChannelTest: true,
  95. ChannelMeta: &relaycommon.ChannelMeta{
  96. HeadersOverride: map[string]any{
  97. "X-Upstream-Trace": "{client_header:X-Trace-Id}",
  98. },
  99. },
  100. }
  101. headers, err := processHeaderOverride(info, ctx)
  102. require.NoError(t, err)
  103. _, ok := headers["X-Upstream-Trace"]
  104. require.False(t, ok)
  105. }
  106. func TestProcessHeaderOverride_NonTestKeepsClientHeaderPlaceholder(t *testing.T) {
  107. t.Parallel()
  108. gin.SetMode(gin.TestMode)
  109. recorder := httptest.NewRecorder()
  110. ctx, _ := gin.CreateTestContext(recorder)
  111. ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
  112. ctx.Request.Header.Set("X-Trace-Id", "trace-123")
  113. info := &relaycommon.RelayInfo{
  114. IsChannelTest: false,
  115. ChannelMeta: &relaycommon.ChannelMeta{
  116. HeadersOverride: map[string]any{
  117. "X-Upstream-Trace": "{client_header:X-Trace-Id}",
  118. },
  119. },
  120. }
  121. headers, err := processHeaderOverride(info, ctx)
  122. require.NoError(t, err)
  123. require.Equal(t, "trace-123", headers["X-Upstream-Trace"])
  124. }
  125. func TestProcessHeaderOverride_PassthroughSkipsAcceptEncoding(t *testing.T) {
  126. t.Parallel()
  127. gin.SetMode(gin.TestMode)
  128. recorder := httptest.NewRecorder()
  129. ctx, _ := gin.CreateTestContext(recorder)
  130. ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
  131. ctx.Request.Header.Set("X-Trace-Id", "trace-123")
  132. ctx.Request.Header.Set("Accept-Encoding", "gzip")
  133. info := &relaycommon.RelayInfo{
  134. IsChannelTest: false,
  135. ChannelMeta: &relaycommon.ChannelMeta{
  136. HeadersOverride: map[string]any{
  137. "*": "",
  138. },
  139. },
  140. }
  141. headers, err := processHeaderOverride(info, ctx)
  142. require.NoError(t, err)
  143. require.Equal(t, "trace-123", headers["X-Trace-Id"])
  144. _, hasAcceptEncoding := headers["Accept-Encoding"]
  145. require.False(t, hasAcceptEncoding)
  146. }