github.com/miolini/go@v0.0.0-20160405192216-fca68c8cb408/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
     6  
     7  import (
     8  	"bytes"
     9  	"net/http"
    10  )
    11  
    12  // ResponseRecorder is an implementation of http.ResponseWriter that
    13  // records its mutations for later inspection in tests.
    14  type ResponseRecorder struct {
    15  	Code      int           // the HTTP response code from WriteHeader
    16  	HeaderMap http.Header   // the HTTP response headers
    17  	Body      *bytes.Buffer // if non-nil, the bytes.Buffer to append written data to
    18  	Flushed   bool
    19  
    20  	stagingMap http.Header // map that handlers manipulate to set headers
    21  	trailerMap http.Header // lazily filled when Trailers() is called
    22  
    23  	wroteHeader bool
    24  }
    25  
    26  // NewRecorder returns an initialized ResponseRecorder.
    27  func NewRecorder() *ResponseRecorder {
    28  	return &ResponseRecorder{
    29  		HeaderMap: make(http.Header),
    30  		Body:      new(bytes.Buffer),
    31  		Code:      200,
    32  	}
    33  }
    34  
    35  // DefaultRemoteAddr is the default remote address to return in RemoteAddr if
    36  // an explicit DefaultRemoteAddr isn't set on ResponseRecorder.
    37  const DefaultRemoteAddr = "1.2.3.4"
    38  
    39  // Header returns the response headers.
    40  func (rw *ResponseRecorder) Header() http.Header {
    41  	m := rw.stagingMap
    42  	if m == nil {
    43  		m = make(http.Header)
    44  		rw.stagingMap = m
    45  	}
    46  	return m
    47  }
    48  
    49  // writeHeader writes a header if it was not written yet and
    50  // detects Content-Type if needed.
    51  //
    52  // bytes or str are the beginning of the response body.
    53  // We pass both to avoid unnecessarily generate garbage
    54  // in rw.WriteString which was created for performance reasons.
    55  // Non-nil bytes win.
    56  func (rw *ResponseRecorder) writeHeader(b []byte, str string) {
    57  	if rw.wroteHeader {
    58  		return
    59  	}
    60  	if len(str) > 512 {
    61  		str = str[:512]
    62  	}
    63  
    64  	m := rw.Header()
    65  
    66  	_, hasType := m["Content-Type"]
    67  	hasTE := m.Get("Transfer-Encoding") != ""
    68  	if !hasType && !hasTE {
    69  		if b == nil {
    70  			b = []byte(str)
    71  		}
    72  		m.Set("Content-Type", http.DetectContentType(b))
    73  	}
    74  
    75  	rw.WriteHeader(200)
    76  }
    77  
    78  // Write always succeeds and writes to rw.Body, if not nil.
    79  func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
    80  	rw.writeHeader(buf, "")
    81  	if rw.Body != nil {
    82  		rw.Body.Write(buf)
    83  	}
    84  	return len(buf), nil
    85  }
    86  
    87  // WriteString always succeeds and writes to rw.Body, if not nil.
    88  func (rw *ResponseRecorder) WriteString(str string) (int, error) {
    89  	rw.writeHeader(nil, str)
    90  	if rw.Body != nil {
    91  		rw.Body.WriteString(str)
    92  	}
    93  	return len(str), nil
    94  }
    95  
    96  // WriteHeader sets rw.Code. After it is called, changing rw.Header
    97  // will not affect rw.HeaderMap.
    98  func (rw *ResponseRecorder) WriteHeader(code int) {
    99  	if rw.wroteHeader {
   100  		return
   101  	}
   102  	rw.Code = code
   103  	rw.wroteHeader = true
   104  	if rw.HeaderMap == nil {
   105  		rw.HeaderMap = make(http.Header)
   106  	}
   107  	for k, vv := range rw.stagingMap {
   108  		vv2 := make([]string, len(vv))
   109  		copy(vv2, vv)
   110  		rw.HeaderMap[k] = vv2
   111  	}
   112  }
   113  
   114  // Flush sets rw.Flushed to true.
   115  func (rw *ResponseRecorder) Flush() {
   116  	if !rw.wroteHeader {
   117  		rw.WriteHeader(200)
   118  	}
   119  	rw.Flushed = true
   120  }
   121  
   122  // Trailers returns any trailers set by the handler. It must be called
   123  // after the handler finished running.
   124  func (rw *ResponseRecorder) Trailers() http.Header {
   125  	if rw.trailerMap != nil {
   126  		return rw.trailerMap
   127  	}
   128  	trailers, ok := rw.HeaderMap["Trailer"]
   129  	if !ok {
   130  		rw.trailerMap = make(http.Header)
   131  		return rw.trailerMap
   132  	}
   133  	rw.trailerMap = make(http.Header, len(trailers))
   134  	for _, k := range trailers {
   135  		switch k {
   136  		case "Transfer-Encoding", "Content-Length", "Trailer":
   137  			// Ignore since forbidden by RFC 2616 14.40.
   138  			continue
   139  		}
   140  		k = http.CanonicalHeaderKey(k)
   141  		vv, ok := rw.stagingMap[k]
   142  		if !ok {
   143  			continue
   144  		}
   145  		vv2 := make([]string, len(vv))
   146  		copy(vv2, vv)
   147  		rw.trailerMap[k] = vv2
   148  	}
   149  	return rw.trailerMap
   150  }