diff --git a/http.go b/http.go index 32d8e46..bcb7bc6 100644 --- a/http.go +++ b/http.go @@ -11,6 +11,7 @@ type HTTPData struct { Res struct { Header http.Header WriteHeader func (statusCode int) + Reset func () } // Req is the HTTP request. Req *http.Request @@ -21,20 +22,16 @@ var _ http.ResponseWriter = new(HTTPResponseRecorder) // HTTPResponseRecorder is an http.ResponseWriter that can buffer a response to // be played back later. type HTTPResponseRecorder struct { - statusCode int - header http.Header - buffer bytes.Buffer + Status int + Head http.Header + buffer bytes.Buffer } func (this *HTTPResponseRecorder) Header () http.Header { - if this.header == nil { - this.header = make(http.Header) + if this.Head == nil { + this.Head = make(http.Header) } - return this.header -} - -func (this *HTTPResponseRecorder) SetHeader (header http.Header) { - this.header = header + return this.Head } func (this *HTTPResponseRecorder) Write (buffer []byte) (int, error) { @@ -42,13 +39,17 @@ func (this *HTTPResponseRecorder) Write (buffer []byte) (int, error) { } func (this *HTTPResponseRecorder) WriteHeader (statusCode int) { - this.statusCode = statusCode + this.Status = statusCode } - // Play replays the response to the given http.ResponseWriter. This resets the // recorder. func (this *HTTPResponseRecorder) Play (res http.ResponseWriter) error { defer this.Reset() + header := res.Header() + for name, value := range this.Head { + header[name] = value + } + res.WriteHeader(this.Status) _, err := io.Copy(res, &this.buffer) return err } @@ -56,6 +57,6 @@ func (this *HTTPResponseRecorder) Play (res http.ResponseWriter) error { // Reset resets this response recorder so it can be used again. func (this *HTTPResponseRecorder) Reset () { this.buffer.Reset() - this.header = nil - this.statusCode = http.StatusOK + this.Head = nil + this.Status = http.StatusOK }