diff --git a/response_writer_test.go b/response_writer_test.go index 259b8fa8..fe4fe25e 100644 --- a/response_writer_test.go +++ b/response_writer_test.go @@ -5,6 +5,8 @@ package gin import ( + "bufio" + "net" "net/http" "net/http/httptest" "testing" @@ -124,6 +126,45 @@ func TestResponseWriterHijack(t *testing.T) { w.Flush() } +type mockHijacker struct { + *httptest.ResponseRecorder + hijacked bool +} + +func (m *mockHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) { + m.hijacked = true + return nil, nil, nil +} + +func TestResponseWriterHijackAfterWrite(t *testing.T) { + // Test case 1: Hijack before writing + hijacker := &mockHijacker{ResponseRecorder: httptest.NewRecorder()} + writer := &responseWriter{} + writer.reset(hijacker) + w := ResponseWriter(writer) + + _, _, err := w.Hijack() + require.NoError(t, err) + assert.True(t, hijacker.hijacked, "Hijack should be called") + assert.True(t, w.Written(), "Written() should be true after Hijack") + + // Test case 2: Hijack after writing + hijacker2 := &mockHijacker{ResponseRecorder: httptest.NewRecorder()} + writer2 := &responseWriter{} + writer2.reset(hijacker2) + w2 := ResponseWriter(writer2) + + _, err = w2.Write([]byte("test")) + require.NoError(t, err) + assert.True(t, w2.Written(), "Written() should be true after Write") + + // Now, try to hijack + _, _, err = w2.Hijack() + require.Error(t, err) + assert.Equal(t, errHijackAlreadyWritten, err, "Hijack after write should return errHijackAlreadyWritten") + assert.False(t, hijacker2.hijacked, "Hijack should not be called after write") +} + func TestResponseWriterFlush(t *testing.T) { testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { writer := &responseWriter{}