package common import ( "bytes" "io" "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) body := `{"model":"gpt-4o-mini","chat_id":"chat_abc","messages":[{"role":"user","content":"hi"}]}` c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(body)) c.Request.Header.Set("Content-Type", "application/json") chatID, err := ExtractTopLevelChatID(c) require.NoError(t, err) require.Equal(t, "chat_abc", chatID) bodyAfter, err := io.ReadAll(c.Request.Body) require.NoError(t, err) require.Equal(t, body, string(bodyAfter)) } 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) } func TestExtractTopLevelChatIDSkipsNonJSONBody(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) body := "--boundary\r\nContent-Disposition: form-data; name=\"chat_id\"\r\n\r\nchat_abc\r\n--boundary--\r\n" c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewBufferString(body)) c.Request.Header.Set("Content-Type", "multipart/form-data; boundary=boundary") chatID, err := ExtractTopLevelChatID(c) require.NoError(t, err) require.Empty(t, chatID) bodyAfter, err := io.ReadAll(c.Request.Body) require.NoError(t, err) require.Equal(t, body, string(bodyAfter)) }