|
5 | 5 | package gin |
6 | 6 |
|
7 | 7 | import ( |
| 8 | + "bufio" |
| 9 | + "net" |
8 | 10 | "net/http" |
9 | 11 | "net/http/httptest" |
10 | 12 | "testing" |
@@ -124,6 +126,74 @@ func TestResponseWriterHijack(t *testing.T) { |
124 | 126 | w.Flush() |
125 | 127 | } |
126 | 128 |
|
| 129 | +type mockHijacker struct { |
| 130 | + *httptest.ResponseRecorder |
| 131 | + hijacked bool |
| 132 | +} |
| 133 | + |
| 134 | +// Hijack implements the http.Hijacker interface. It just records that it was called. |
| 135 | +func (m *mockHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) { |
| 136 | + m.hijacked = true |
| 137 | + return nil, nil, nil |
| 138 | +} |
| 139 | + |
| 140 | +func TestResponseWriterHijackAfterWrite(t *testing.T) { |
| 141 | + tests := []struct { |
| 142 | + name string |
| 143 | + action func(w ResponseWriter) error // Action to perform before hijacking |
| 144 | + expectWrittenBeforeHijack bool |
| 145 | + expectHijackSuccess bool |
| 146 | + expectWrittenAfterHijack bool |
| 147 | + expectError error |
| 148 | + }{ |
| 149 | + { |
| 150 | + name: "hijack before write should succeed", |
| 151 | + action: func(w ResponseWriter) error { return nil }, |
| 152 | + expectWrittenBeforeHijack: false, |
| 153 | + expectHijackSuccess: true, |
| 154 | + expectWrittenAfterHijack: true, // Hijack itself marks the writer as written |
| 155 | + expectError: nil, |
| 156 | + }, |
| 157 | + { |
| 158 | + name: "hijack after write should fail", |
| 159 | + action: func(w ResponseWriter) error { |
| 160 | + _, err := w.Write([]byte("test")) |
| 161 | + return err |
| 162 | + }, |
| 163 | + expectWrittenBeforeHijack: true, |
| 164 | + expectHijackSuccess: false, |
| 165 | + expectWrittenAfterHijack: true, |
| 166 | + expectError: errHijackAlreadyWritten, |
| 167 | + }, |
| 168 | + } |
| 169 | + |
| 170 | + for _, tc := range tests { |
| 171 | + t.Run(tc.name, func(t *testing.T) { |
| 172 | + hijacker := &mockHijacker{ResponseRecorder: httptest.NewRecorder()} |
| 173 | + writer := &responseWriter{} |
| 174 | + writer.reset(hijacker) |
| 175 | + w := ResponseWriter(writer) |
| 176 | + |
| 177 | + // Check initial state |
| 178 | + assert.False(t, w.Written(), "should not be written initially") |
| 179 | + |
| 180 | + // Perform pre-hijack action |
| 181 | + require.NoError(t, tc.action(w), "unexpected error during pre-hijack action") |
| 182 | + |
| 183 | + // Check state before hijacking |
| 184 | + assert.Equal(t, tc.expectWrittenBeforeHijack, w.Written(), "unexpected w.Written() state before hijack") |
| 185 | + |
| 186 | + // Attempt to hijack |
| 187 | + _, _, hijackErr := w.Hijack() |
| 188 | + |
| 189 | + // Check results |
| 190 | + require.ErrorIs(t, hijackErr, tc.expectError, "unexpected error from Hijack()") |
| 191 | + assert.Equal(t, tc.expectHijackSuccess, hijacker.hijacked, "unexpected hijacker.hijacked state") |
| 192 | + assert.Equal(t, tc.expectWrittenAfterHijack, w.Written(), "unexpected w.Written() state after hijack") |
| 193 | + }) |
| 194 | + } |
| 195 | +} |
| 196 | + |
127 | 197 | func TestResponseWriterFlush(t *testing.T) { |
128 | 198 | testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
129 | 199 | writer := &responseWriter{} |
|
0 commit comments