Nie możesz wybrać więcej, niż 25 tematów Tematy muszą się zaczynać od litery lub cyfry, mogą zawierać myślniki ('-') i mogą mieć do 35 znaków.
 
 
 

90 wiersze
2.5 KiB

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