diff --git a/rest/handler/timeouthandler.go b/rest/handler/timeouthandler.go index ec34180a..52d82a80 100644 --- a/rest/handler/timeouthandler.go +++ b/rest/handler/timeouthandler.go @@ -127,6 +127,12 @@ type timeoutWriter struct { var _ http.Pusher = (*timeoutWriter)(nil) +func (tw *timeoutWriter) Flush() { + if flusher, ok := tw.w.(http.Flusher); ok { + flusher.Flush() + } +} + func (tw *timeoutWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { if hijacked, ok := tw.w.(http.Hijacker); ok { return hijacked.Hijack() diff --git a/rest/handler/timeouthandler_test.go b/rest/handler/timeouthandler_test.go index 3fd8bdd2..acbda030 100644 --- a/rest/handler/timeouthandler_test.go +++ b/rest/handler/timeouthandler_test.go @@ -164,6 +164,24 @@ func TestTimeoutHijack(t *testing.T) { }) } +func TestTimeoutFlush(t *testing.T) { + timeoutHandler := TimeoutHandler(time.Minute) + handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + flusher, ok := w.(http.Flusher) + if !ok { + http.Error(w, "Streaming unsupported!", http.StatusInternalServerError) + return + } + + flusher.Flush() + })) + + req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody) + resp := httptest.NewRecorder() + handler.ServeHTTP(resp, req) + assert.Equal(t, http.StatusOK, resp.Code) +} + func TestTimeoutPusher(t *testing.T) { handler := &timeoutWriter{ w: mockedPusher{},