mirror of
https://github.com/zeromicro/go-zero.git
synced 2025-01-24 18:00:24 +08:00
e97e1f10db
* simplify code with http.Flusher type conversion * simplify code with http.Flusher type conversion, better version
34 lines
880 B
Go
34 lines
880 B
Go
package security
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func TestWithCodeResponseWriter(t *testing.T) {
|
|
req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
|
|
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
cw := &WithCodeResponseWriter{Writer: w}
|
|
|
|
cw.Header().Set("X-Test", "test")
|
|
cw.WriteHeader(http.StatusServiceUnavailable)
|
|
assert.Equal(t, cw.Code, http.StatusServiceUnavailable)
|
|
|
|
_, err := cw.Write([]byte("content"))
|
|
assert.Nil(t, err)
|
|
|
|
flusher, ok := http.ResponseWriter(cw).(http.Flusher)
|
|
assert.True(t, ok)
|
|
flusher.Flush()
|
|
})
|
|
|
|
resp := httptest.NewRecorder()
|
|
handler.ServeHTTP(resp, req)
|
|
assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
|
|
assert.Equal(t, "test", resp.Header().Get("X-Test"))
|
|
assert.Equal(t, "content", resp.Body.String())
|
|
}
|