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) }