Ver código fonte

timeout writer add hijack

yangjinheng 2 anos atrás
pai
commit
dbc8f9faca
2 arquivos alterados com 27 adições e 0 exclusões
  1. 2 0
      rest/handler/timeouthandler.go
  2. 25 0
      rest/handler/timeouthandler_test.go

+ 2 - 0
rest/handler/timeouthandler.go

@@ -1,11 +1,13 @@
 package handler
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"errors"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"path"
 	"runtime"

+ 25 - 0
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{},