github.com/sean-/go@v0.0.0-20151219100004-97f854cd7bb6/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  		if rw.HeaderMap == nil {
    69  			rw.HeaderMap = make(http.Header)
    70  		}
    71  		rw.HeaderMap.Set("Content-Type", http.DetectContentType(b))
    72  	}
    73  
    74  	rw.WriteHeader(200)
    75  }
    76  
    77  // Write always succeeds and writes to rw.Body, if not nil.
    78  func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
    79  	rw.writeHeader(buf, "")
    80  	if rw.Body != nil {
    81  		rw.Body.Write(buf)
    82  	}
    83  	return len(buf), nil
    84  }
    85  
    86  // WriteString always succeeds and writes to rw.Body, if not nil.
    87  func (rw *ResponseRecorder) WriteString(str string) (int, error) {
    88  	rw.writeHeader(nil, str)
    89  	if rw.Body != nil {
    90  		rw.Body.WriteString(str)
    91  	}
    92  	return len(str), nil
    93  }
    94  
    95  // WriteHeader sets rw.Code.
    96  func (rw *ResponseRecorder) WriteHeader(code int) {
    97  	if !rw.wroteHeader {
    98  		rw.Code = code
    99  		rw.wroteHeader = true
   100  	}
   101  }
   102  
   103  // Flush sets rw.Flushed to true.
   104  func (rw *ResponseRecorder) Flush() {
   105  	if !rw.wroteHeader {
   106  		rw.WriteHeader(200)
   107  	}
   108  	rw.Flushed = true
   109  }