|
- package common
-
- import (
- "bytes"
- "net/http"
- "net/http/httptest"
- "testing"
-
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/require"
- )
-
- func TestExtractChatIDFromReusableBody(t *testing.T) {
- gin.SetMode(gin.TestMode)
- w := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(w)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"gpt-4o-mini","chat_id":"chat_abc","messages":[{"role":"user","content":"hi"}]}`))
- c.Request.Header.Set("Content-Type", "application/json")
-
- chatID, err := ExtractTopLevelChatID(c)
- require.NoError(t, err)
- require.Equal(t, "chat_abc", chatID)
- }
-
- func TestExtractChatIDReturnsEmptyWhenMissing(t *testing.T) {
- gin.SetMode(gin.TestMode)
- w := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(w)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`))
- c.Request.Header.Set("Content-Type", "application/json")
-
- chatID, err := ExtractTopLevelChatID(c)
- require.NoError(t, err)
- require.Empty(t, chatID)
- }
-
- func TestCaptureUpstreamIDFromHeaderPrefersXRequestID(t *testing.T) {
- info := &RelayInfo{}
- resp := &http.Response{
- Header: http.Header{
- "X-Request-Id": []string{"xreq_123"},
- "Request-Id": []string{"req_456"},
- },
- }
-
- CaptureUpstreamIDFromHTTPResponse(info, resp)
- require.Equal(t, "xreq_123", info.UpstreamID)
- }
-
- func TestCaptureUpstreamIDFallsBackToRequestID(t *testing.T) {
- info := &RelayInfo{}
- resp := &http.Response{
- Header: http.Header{
- "Request-Id": []string{"req_456"},
- },
- }
-
- CaptureUpstreamIDFromHTTPResponse(info, resp)
- require.Equal(t, "req_456", info.UpstreamID)
- }
-
- func TestSetAndGetChatIDFromGinContext(t *testing.T) {
- gin.SetMode(gin.TestMode)
- w := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(w)
-
- SetRelayChatID(c, "chat_ctx_1")
- require.Equal(t, "chat_ctx_1", GetRelayChatID(c))
- }
-
- func TestCaptureUpstreamIDDoesNotOverwriteExistingValueWithEmptyHeader(t *testing.T) {
- info := &RelayInfo{UpstreamID: "existing_upstream"}
- resp := &http.Response{Header: http.Header{}}
-
- CaptureUpstreamIDFromHTTPResponse(info, resp)
- require.Equal(t, "existing_upstream", info.UpstreamID)
- }
-
- func TestCaptureUpstreamIDFromDialResponseHeaders(t *testing.T) {
- info := &RelayInfo{}
- resp := &http.Response{
- Header: http.Header{
- "X-Request-Id": []string{"ws_upstream_1"},
- },
- }
-
- CaptureUpstreamIDFromHTTPResponse(info, resp)
- require.Equal(t, "ws_upstream_1", info.UpstreamID)
- }
|