github.com/mdempsky/go@v0.0.0-20151201204031-5dd372bd1e70/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  // writeHeader writes a header if it was not written yet and
    48  // detects Content-Type if needed.
    49  //
    50  // bytes or str are the beginning of the response body.
    51  // We pass both to avoid unnecessarily generate garbage
    52  // in rw.WriteString which was created for performance reasons.
    53  // Non-nil bytes win.
    54  func (rw *ResponseRecorder) writeHeader(b []byte, str string) {
    55  	if rw.wroteHeader {
    56  		return
    57  	}
    58  	if len(str) > 512 {
    59  		str = str[:512]
    60  	}
    61  
    62  	_, hasType := rw.HeaderMap["Content-Type"]
    63  	hasTE := rw.HeaderMap.Get("Transfer-Encoding") != ""
    64  	if !hasType && !hasTE {
    65  		if b == nil {
    66  			b = []byte(str)
    67  		}
    68  		rw.HeaderMap.Set("Content-Type", http.DetectContentType(b))
    69  	}
    70  
    71  	rw.WriteHeader(200)
    72  }
    73  
    74  // Write always succeeds and writes to rw.Body, if not nil.
    75  func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
    76  	rw.writeHeader(buf, "")
    77  	if rw.Body != nil {
    78  		rw.Body.Write(buf)
    79  	}
    80  	return len(buf), nil
    81  }
    82  
    83  // WriteString always succeeds and writes to rw.Body, if not nil.
    84  func (rw *ResponseRecorder) WriteString(str string) (int, error) {
    85  	rw.writeHeader(nil, str)
    86  	if rw.Body != nil {
    87  		rw.Body.WriteString(str)
    88  	}
    89  	return len(str), nil
    90  }
    91  
    92  // WriteHeader sets rw.Code.
    93  func (rw *ResponseRecorder) WriteHeader(code int) {
    94  	if !rw.wroteHeader {
    95  		rw.Code = code
    96  		rw.wroteHeader = true
    97  	}
    98  }
    99  
   100  // Flush sets rw.Flushed to true.
   101  func (rw *ResponseRecorder) Flush() {
   102  	if !rw.wroteHeader {
   103  		rw.WriteHeader(200)
   104  	}
   105  	rw.Flushed = true
   106  }