fix review

response_writer_test.go

@@ -131,39 +131,54 @@ type mockHijacker struct {
 	hijacked bool
 }
 
+// Hijack implements the http.Hijacker interface. It just records that it was called.
 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
+	tests := []struct {
-	hijacker := &mockHijacker{ResponseRecorder: httptest.NewRecorder()}
+		name          string
-	writer := &responseWriter{}
+		action        func(w ResponseWriter) error // Action to perform before hijacking
-	writer.reset(hijacker)
+		expectHijack  bool
-	w := ResponseWriter(writer)
+		expectWritten bool
+		expectError   error
+	}{
+		{
+			name:          "hijack before write should succeed",
+			action:        func(w ResponseWriter) error { return nil },
+			expectHijack:  true,
+			expectWritten: true, // Hijack itself marks the writer as written
+			expectError:   nil,
+		},
+		{
+			name: "hijack after write should fail",
+			action: func(w ResponseWriter) error {
+				_, err := w.Write([]byte("test"))
+				return err
+			},
+			expectHijack:  false,
+			expectWritten: true,
+			expectError:   errHijackAlreadyWritten,
+		},
+	}
 
-	_, _, err := w.Hijack()
+	for _, tc := range tests {
-	require.NoError(t, err)
+		t.Run(tc.name, func(t *testing.T) {
-	assert.True(t, hijacker.hijacked, "Hijack should be called")
+			hijacker := &mockHijacker{ResponseRecorder: httptest.NewRecorder()}
-	assert.True(t, w.Written(), "Written() should be true after Hijack")
+			writer := &responseWriter{}
+			writer.reset(hijacker)
+			w := ResponseWriter(writer)
 
-	// Test case 2: Hijack after writing
+			require.NoError(t, tc.action(w), "unexpected error during pre-hijack action")
-	hijacker2 := &mockHijacker{ResponseRecorder: httptest.NewRecorder()}
-	writer2 := &responseWriter{}
-	writer2.reset(hijacker2)
-	w2 := ResponseWriter(writer2)
 
-	_, err = w2.Write([]byte("test"))
+			_, _, hijackErr := w.Hijack()
-	require.NoError(t, err)
+			assert.ErrorIs(t, hijackErr, tc.expectError, "unexpected error from Hijack()")
-	assert.True(t, w2.Written(), "Written() should be true after Write")
+			assert.Equal(t, tc.expectHijack, hijacker.hijacked, "unexpected hijacker.hijacked state")
-
+			assert.Equal(t, tc.expectWritten, w.Written(), "unexpected w.Written() state")
-	// 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) {