github.com/yanyiwu/go@v0.0.0-20150106053140-03d6637dbb7f/src/net/http/httptest/recorder.go (about)

     1  // Copyright 2011 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package httptest provides utilities for HTTP testing.
     6  package httptest
     7  
     8  import (
     9  	"bytes"
    10  	"net/http"
    11  )
    12  
    13  // ResponseRecorder is an implementation of http.ResponseWriter that
    14  // records its mutations for later inspection in tests.
    15  type ResponseRecorder struct {
    16  	Code      int           // the HTTP response code from WriteHeader
    17  	HeaderMap http.Header   // the HTTP response headers
    18  	Body      *bytes.Buffer // if non-nil, the bytes.Buffer to append written data to
    19  	Flushed   bool
    20  
    21  	wroteHeader bool
    22  }
    23  
    24  // NewRecorder returns an initialized ResponseRecorder.
    25  func NewRecorder() *ResponseRecorder {
    26  	return &ResponseRecorder{
    27  		HeaderMap: make(http.Header),
    28  		Body:      new(bytes.Buffer),
    29  		Code:      200,
    30  	}
    31  }
    32  
    33  // DefaultRemoteAddr is the default remote address to return in RemoteAddr if
    34  // an explicit DefaultRemoteAddr isn't set on ResponseRecorder.
    35  const DefaultRemoteAddr = "1.2.3.4"
    36  
    37  // Header returns the response headers.
    38  func (rw *ResponseRecorder) Header() http.Header {
    39  	m := rw.HeaderMap
    40  	if m == nil {
    41  		m = make(http.Header)
    42  		rw.HeaderMap = m
    43  	}
    44  	return m
    45  }
    46  
    47  // Write always succeeds and writes to rw.Body, if not nil.
    48  func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
    49  	if !rw.wroteHeader {
    50  		rw.WriteHeader(200)
    51  	}
    52  	if rw.Body != nil {
    53  		rw.Body.Write(buf)
    54  	}
    55  	return len(buf), nil
    56  }
    57  
    58  // WriteHeader sets rw.Code.
    59  func (rw *ResponseRecorder) WriteHeader(code int) {
    60  	if !rw.wroteHeader {
    61  		rw.Code = code
    62  	}
    63  	rw.wroteHeader = true
    64  }
    65  
    66  // Flush sets rw.Flushed to true.
    67  func (rw *ResponseRecorder) Flush() {
    68  	if !rw.wroteHeader {
    69  		rw.WriteHeader(200)
    70  	}
    71  	rw.Flushed = true
    72  }