github.com/lingyao2333/mo-zero@v1.4.1/rest/internal/response/headeronceresponsewriter_test.go (about)

     1  package response
     2  
     3  import (
     4  	"bufio"
     5  	"net"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  func TestHeaderOnceResponseWriter_Flush(t *testing.T) {
    14  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    15  	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    16  		cw := NewHeaderOnceResponseWriter(w)
    17  		cw.Header().Set("X-Test", "test")
    18  		cw.WriteHeader(http.StatusServiceUnavailable)
    19  		cw.WriteHeader(http.StatusExpectationFailed)
    20  		_, err := cw.Write([]byte("content"))
    21  		assert.Nil(t, err)
    22  
    23  		flusher, ok := cw.(http.Flusher)
    24  		assert.True(t, ok)
    25  		flusher.Flush()
    26  	})
    27  
    28  	resp := httptest.NewRecorder()
    29  	handler.ServeHTTP(resp, req)
    30  	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
    31  	assert.Equal(t, "test", resp.Header().Get("X-Test"))
    32  	assert.Equal(t, "content", resp.Body.String())
    33  }
    34  
    35  func TestHeaderOnceResponseWriter_Hijack(t *testing.T) {
    36  	resp := httptest.NewRecorder()
    37  	writer := &HeaderOnceResponseWriter{
    38  		w: resp,
    39  	}
    40  	assert.NotPanics(t, func() {
    41  		writer.Hijack()
    42  	})
    43  
    44  	writer = &HeaderOnceResponseWriter{
    45  		w: mockedHijackable{resp},
    46  	}
    47  	assert.NotPanics(t, func() {
    48  		writer.Hijack()
    49  	})
    50  }
    51  
    52  type mockedHijackable struct {
    53  	*httptest.ResponseRecorder
    54  }
    55  
    56  func (m mockedHijackable) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    57  	return nil, nil, nil
    58  }