diff --git a/rest/handler/timeouthandler.go b/rest/handler/timeouthandler.go index fb686285..0becace3 100644 --- a/rest/handler/timeouthandler.go +++ b/rest/handler/timeouthandler.go @@ -1,11 +1,13 @@ package handler import ( + "bufio" "bytes" "context" "errors" "fmt" "io" + "net" "net/http" "path" "runtime" diff --git a/rest/handler/timeouthandler_test.go b/rest/handler/timeouthandler_test.go index 73e75cce..23eb88c6 100644 --- a/rest/handler/timeouthandler_test.go +++ b/rest/handler/timeouthandler_test.go @@ -10,6 +10,7 @@ import ( "time" "github.com/stretchr/testify/assert" + "github.com/zeromicro/go-zero/rest/internal/response" ) func init() { @@ -134,6 +135,30 @@ func TestTimeoutClientClosed(t *testing.T) { assert.Equal(t, statusClientClosedRequest, resp.Code) } +func TestTimeoutHijack(t *testing.T) { + resp := httptest.NewRecorder() + + writer := &timeoutWriter{ + w: &response.WithCodeResponseWriter{ + Writer: resp, + }, + } + + assert.NotPanics(t, func() { + writer.Hijack() + }) + + writer = &timeoutWriter{ + w: &response.WithCodeResponseWriter{ + Writer: mockedHijackable{resp}, + }, + } + + assert.NotPanics(t, func() { + writer.Hijack() + }) +} + func TestTimeoutPusher(t *testing.T) { handler := &timeoutWriter{ w: mockedPusher{},