diff --git a/fsttest/recorder.go b/fsttest/recorder.go index 7e4a60e..2648b75 100644 --- a/fsttest/recorder.go +++ b/fsttest/recorder.go @@ -12,9 +12,10 @@ import ( // ResponseRecorder is an implementation of fsthttp.ResponseWriter that // records its mutations for later inspection in tests. type ResponseRecorder struct { - Code int - HeaderMap fsthttp.Header - Body *bytes.Buffer + Code int + HeaderMap fsthttp.Header + Body *bytes.Buffer + headersSent bool } // NewRecorder returns an initialized ResponseRecorder. @@ -28,12 +29,19 @@ func NewRecorder() *ResponseRecorder { // Header returns the response headers to mutate within a handler. func (r *ResponseRecorder) Header() fsthttp.Header { - return r.HeaderMap + if !r.headersSent { + return r.HeaderMap + } + // Once sent, return a copy so any changes are discarded. + return r.HeaderMap.Clone() } // WriteHeader records the response code. func (r *ResponseRecorder) WriteHeader(code int) { - r.Code = code + if !r.headersSent { + r.Code = code + r.headersSent = true + } } // Write records the response body. The data is written to the Body