recorder.go

Documentation: net/http/httptest

     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  	"fmt"
    10  	"io/ioutil"
    11  	"net/http"
    12  	"strconv"
    13  	"strings"
    14  
    15  	"internal/x/net/http/httpguts"
    16  )
    17  
    18  // ResponseRecorder is an implementation of http.ResponseWriter that
    19  // records its mutations for later inspection in tests.
    20  type ResponseRecorder struct {
    21  	// Code is the HTTP response code set by WriteHeader.
    22  	//
    23  	// Note that if a Handler never calls WriteHeader or Write,
    24  	// this might end up being 0, rather than the implicit
    25  	// http.StatusOK. To get the implicit value, use the Result
    26  	// method.
    27  	Code int
    28  
    29  	// HeaderMap contains the headers explicitly set by the Handler.
    30  	// It is an internal detail.
    31  	//
    32  	// Deprecated: HeaderMap exists for historical compatibility
    33  	// and should not be used. To access the headers returned by a handler,
    34  	// use the Response.Header map as returned by the Result method.
    35  	HeaderMap http.Header
    36  
    37  	// Body is the buffer to which the Handler's Write calls are sent.
    38  	// If nil, the Writes are silently discarded.
    39  	Body *bytes.Buffer
    40  
    41  	// Flushed is whether the Handler called Flush.
    42  	Flushed bool
    43  
    44  	result      *http.Response // cache of Result's return value
    45  	snapHeader  http.Header    // snapshot of HeaderMap at first Write
    46  	wroteHeader bool
    47  }
    48  
    49  // NewRecorder returns an initialized ResponseRecorder.
    50  func NewRecorder() *ResponseRecorder {
    51  	return &ResponseRecorder{
    52  		HeaderMap: make(http.Header),
    53  		Body:      new(bytes.Buffer),
    54  		Code:      200,
    55  	}
    56  }
    57  
    58  // DefaultRemoteAddr is the default remote address to return in RemoteAddr if
    59  // an explicit DefaultRemoteAddr isn't set on ResponseRecorder.
    60  const DefaultRemoteAddr = "1.2.3.4"
    61  
    62  // Header returns the response headers.
    63  func (rw *ResponseRecorder) Header() http.Header {
    64  	m := rw.HeaderMap
    65  	if m == nil {
    66  		m = make(http.Header)
    67  		rw.HeaderMap = m
    68  	}
    69  	return m
    70  }
    71  
    72  // writeHeader writes a header if it was not written yet and
    73  // detects Content-Type if needed.
    74  //
    75  // bytes or str are the beginning of the response body.
    76  // We pass both to avoid unnecessarily generate garbage
    77  // in rw.WriteString which was created for performance reasons.
    78  // Non-nil bytes win.
    79  func (rw *ResponseRecorder) writeHeader(b []byte, str string) {
    80  	if rw.wroteHeader {
    81  		return
    82  	}
    83  	if len(str) > 512 {
    84  		str = str[:512]
    85  	}
    86  
    87  	m := rw.Header()
    88  
    89  	_, hasType := m["Content-Type"]
    90  	hasTE := m.Get("Transfer-Encoding") != ""
    91  	if !hasType && !hasTE {
    92  		if b == nil {
    93  			b = []byte(str)
    94  		}
    95  		m.Set("Content-Type", http.DetectContentType(b))
    96  	}
    97  
    98  	rw.WriteHeader(200)
    99  }
   100  
   101  // Write always succeeds and writes to rw.Body, if not nil.
   102  func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
   103  	rw.writeHeader(buf, "")
   104  	if rw.Body != nil {
   105  		rw.Body.Write(buf)
   106  	}
   107  	return len(buf), nil
   108  }
   109  
   110  // WriteString always succeeds and writes to rw.Body, if not nil.
   111  func (rw *ResponseRecorder) WriteString(str string) (int, error) {
   112  	rw.writeHeader(nil, str)
   113  	if rw.Body != nil {
   114  		rw.Body.WriteString(str)
   115  	}
   116  	return len(str), nil
   117  }
   118  
   119  // WriteHeader sets rw.Code. After it is called, changing rw.Header
   120  // will not affect rw.HeaderMap.
   121  func (rw *ResponseRecorder) WriteHeader(code int) {
   122  	if rw.wroteHeader {
   123  		return
   124  	}
   125  	rw.Code = code
   126  	rw.wroteHeader = true
   127  	if rw.HeaderMap == nil {
   128  		rw.HeaderMap = make(http.Header)
   129  	}
   130  	rw.snapHeader = cloneHeader(rw.HeaderMap)
   131  }
   132  
   133  func cloneHeader(h http.Header) http.Header {
   134  	h2 := make(http.Header, len(h))
   135  	for k, vv := range h {
   136  		vv2 := make([]string, len(vv))
   137  		copy(vv2, vv)
   138  		h2[k] = vv2
   139  	}
   140  	return h2
   141  }
   142  
   143  // Flush sets rw.Flushed to true.
   144  func (rw *ResponseRecorder) Flush() {
   145  	if !rw.wroteHeader {
   146  		rw.WriteHeader(200)
   147  	}
   148  	rw.Flushed = true
   149  }
   150  
   151  // Result returns the response generated by the handler.
   152  //
   153  // The returned Response will have at least its StatusCode,
   154  // Header, Body, and optionally Trailer populated.
   155  // More fields may be populated in the future, so callers should
   156  // not DeepEqual the result in tests.
   157  //
   158  // The Response.Header is a snapshot of the headers at the time of the
   159  // first write call, or at the time of this call, if the handler never
   160  // did a write.
   161  //
   162  // The Response.Body is guaranteed to be non-nil and Body.Read call is
   163  // guaranteed to not return any error other than io.EOF.
   164  //
   165  // Result must only be called after the handler has finished running.
   166  func (rw *ResponseRecorder) Result() *http.Response {
   167  	if rw.result != nil {
   168  		return rw.result
   169  	}
   170  	if rw.snapHeader == nil {
   171  		rw.snapHeader = cloneHeader(rw.HeaderMap)
   172  	}
   173  	res := &http.Response{
   174  		Proto:      "HTTP/1.1",
   175  		ProtoMajor: 1,
   176  		ProtoMinor: 1,
   177  		StatusCode: rw.Code,
   178  		Header:     rw.snapHeader,
   179  	}
   180  	rw.result = res
   181  	if res.StatusCode == 0 {
   182  		res.StatusCode = 200
   183  	}
   184  	res.Status = fmt.Sprintf("%03d %s", res.StatusCode, http.StatusText(res.StatusCode))
   185  	if rw.Body != nil {
   186  		res.Body = ioutil.NopCloser(bytes.NewReader(rw.Body.Bytes()))
   187  	} else {
   188  		res.Body = http.NoBody
   189  	}
   190  	res.ContentLength = parseContentLength(res.Header.Get("Content-Length"))
   191  
   192  	if trailers, ok := rw.snapHeader["Trailer"]; ok {
   193  		res.Trailer = make(http.Header, len(trailers))
   194  		for _, k := range trailers {
   195  			k = http.CanonicalHeaderKey(k)
   196  			if !httpguts.ValidTrailerHeader(k) {
   197  				// Ignore since forbidden by RFC 7230, section 4.1.2.
   198  				continue
   199  			}
   200  			vv, ok := rw.HeaderMap[k]
   201  			if !ok {
   202  				continue
   203  			}
   204  			vv2 := make([]string, len(vv))
   205  			copy(vv2, vv)
   206  			res.Trailer[k] = vv2
   207  		}
   208  	}
   209  	for k, vv := range rw.HeaderMap {
   210  		if !strings.HasPrefix(k, http.TrailerPrefix) {
   211  			continue
   212  		}
   213  		if res.Trailer == nil {
   214  			res.Trailer = make(http.Header)
   215  		}
   216  		for _, v := range vv {
   217  			res.Trailer.Add(strings.TrimPrefix(k, http.TrailerPrefix), v)
   218  		}
   219  	}
   220  	return res
   221  }
   222  
   223  // parseContentLength trims whitespace from s and returns -1 if no value
   224  // is set, or the value if it's >= 0.
   225  //
   226  // This a modified version of same function found in net/http/transfer.go. This
   227  // one just ignores an invalid header.
   228  func parseContentLength(cl string) int64 {
   229  	cl = strings.TrimSpace(cl)
   230  	if cl == "" {
   231  		return -1
   232  	}
   233  	n, err := strconv.ParseInt(cl, 10, 64)
   234  	if err != nil {
   235  		return -1
   236  	}
   237  	return n
   238  }
   239
View as plain text