25'ten fazla konu seçemezsiniz Konular bir harf veya rakamla başlamalı, kısa çizgiler ('-') içerebilir ve en fazla 35 karakter uzunluğunda olabilir.
 
 
 

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