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 }