Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

113 строки
3.3 KiB

  1. package common
  2. import (
  3. "bytes"
  4. "io"
  5. "net/http"
  6. "net/http/httptest"
  7. "testing"
  8. "github.com/gin-gonic/gin"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestExtractChatIDFromReusableBody(t *testing.T) {
  12. gin.SetMode(gin.TestMode)
  13. w := httptest.NewRecorder()
  14. c, _ := gin.CreateTestContext(w)
  15. body := `{"model":"gpt-4o-mini","chat_id":"chat_abc","messages":[{"role":"user","content":"hi"}]}`
  16. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(body))
  17. c.Request.Header.Set("Content-Type", "application/json")
  18. chatID, err := ExtractTopLevelChatID(c)
  19. require.NoError(t, err)
  20. require.Equal(t, "chat_abc", chatID)
  21. bodyAfter, err := io.ReadAll(c.Request.Body)
  22. require.NoError(t, err)
  23. require.Equal(t, body, string(bodyAfter))
  24. }
  25. func TestExtractChatIDReturnsEmptyWhenMissing(t *testing.T) {
  26. gin.SetMode(gin.TestMode)
  27. w := httptest.NewRecorder()
  28. c, _ := gin.CreateTestContext(w)
  29. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`))
  30. c.Request.Header.Set("Content-Type", "application/json")
  31. chatID, err := ExtractTopLevelChatID(c)
  32. require.NoError(t, err)
  33. require.Empty(t, chatID)
  34. }
  35. func TestCaptureUpstreamIDFromHeaderPrefersXRequestID(t *testing.T) {
  36. info := &RelayInfo{}
  37. resp := &http.Response{
  38. Header: http.Header{
  39. "X-Request-Id": []string{"xreq_123"},
  40. "Request-Id": []string{"req_456"},
  41. },
  42. }
  43. CaptureUpstreamIDFromHTTPResponse(info, resp)
  44. require.Equal(t, "xreq_123", info.UpstreamID)
  45. }
  46. func TestCaptureUpstreamIDFallsBackToRequestID(t *testing.T) {
  47. info := &RelayInfo{}
  48. resp := &http.Response{
  49. Header: http.Header{
  50. "Request-Id": []string{"req_456"},
  51. },
  52. }
  53. CaptureUpstreamIDFromHTTPResponse(info, resp)
  54. require.Equal(t, "req_456", info.UpstreamID)
  55. }
  56. func TestSetAndGetChatIDFromGinContext(t *testing.T) {
  57. gin.SetMode(gin.TestMode)
  58. w := httptest.NewRecorder()
  59. c, _ := gin.CreateTestContext(w)
  60. SetRelayChatID(c, "chat_ctx_1")
  61. require.Equal(t, "chat_ctx_1", GetRelayChatID(c))
  62. }
  63. func TestCaptureUpstreamIDDoesNotOverwriteExistingValueWithEmptyHeader(t *testing.T) {
  64. info := &RelayInfo{UpstreamID: "existing_upstream"}
  65. resp := &http.Response{Header: http.Header{}}
  66. CaptureUpstreamIDFromHTTPResponse(info, resp)
  67. require.Equal(t, "existing_upstream", info.UpstreamID)
  68. }
  69. func TestCaptureUpstreamIDFromDialResponseHeaders(t *testing.T) {
  70. info := &RelayInfo{}
  71. resp := &http.Response{
  72. Header: http.Header{
  73. "X-Request-Id": []string{"ws_upstream_1"},
  74. },
  75. }
  76. CaptureUpstreamIDFromHTTPResponse(info, resp)
  77. require.Equal(t, "ws_upstream_1", info.UpstreamID)
  78. }
  79. func TestExtractTopLevelChatIDSkipsNonJSONBody(t *testing.T) {
  80. gin.SetMode(gin.TestMode)
  81. w := httptest.NewRecorder()
  82. c, _ := gin.CreateTestContext(w)
  83. body := "--boundary\r\nContent-Disposition: form-data; name=\"chat_id\"\r\n\r\nchat_abc\r\n--boundary--\r\n"
  84. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewBufferString(body))
  85. c.Request.Header.Set("Content-Type", "multipart/form-data; boundary=boundary")
  86. chatID, err := ExtractTopLevelChatID(c)
  87. require.NoError(t, err)
  88. require.Empty(t, chatID)
  89. bodyAfter, err := io.ReadAll(c.Request.Body)
  90. require.NoError(t, err)
  91. require.Equal(t, body, string(bodyAfter))
  92. }