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

     1  package response
     2  
     3  import (
     4  	"net/http"
     5  	"net/http/httptest"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestWithCodeResponseWriter(t *testing.T) {
    12  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    13  	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    14  		cw := &WithCodeResponseWriter{Writer: w}
    15  
    16  		cw.Header().Set("X-Test", "test")
    17  		cw.WriteHeader(http.StatusServiceUnavailable)
    18  		assert.Equal(t, cw.Code, http.StatusServiceUnavailable)
    19  
    20  		_, err := cw.Write([]byte("content"))
    21  		assert.Nil(t, err)
    22  
    23  		flusher, ok := http.ResponseWriter(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 TestWithCodeResponseWriter_Hijack(t *testing.T) {
    36  	resp := httptest.NewRecorder()
    37  	writer := &WithCodeResponseWriter{
    38  		Writer: resp,
    39  	}
    40  	assert.NotPanics(t, func() {
    41  		writer.Hijack()
    42  	})
    43  
    44  	writer = &WithCodeResponseWriter{
    45  		Writer: mockedHijackable{resp},
    46  	}
    47  	assert.NotPanics(t, func() {
    48  		writer.Hijack()
    49  	})
    50  }