github.com/naoina/kocha@v0.7.1-0.20171129072645-78c7a531f799/response.go (about)

     1  package kocha
     2  
     3  import (
     4  	"io"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"sync"
     8  )
     9  
    10  var (
    11  	_ http.ResponseWriter = &Response{}
    12  
    13  	responsePool = &sync.Pool{
    14  		New: func() interface{} {
    15  			return &Response{}
    16  		},
    17  	}
    18  )
    19  
    20  // Response represents a response.
    21  type Response struct {
    22  	http.ResponseWriter
    23  
    24  	ContentType string
    25  	StatusCode  int
    26  
    27  	cookies []*http.Cookie
    28  	resp    *httptest.ResponseRecorder
    29  }
    30  
    31  // newResponse returns a new Response that responds to rw.
    32  func newResponse() *Response {
    33  	r := responsePool.Get().(*Response)
    34  	r.reset()
    35  	r.ContentType = ""
    36  	r.cookies = r.cookies[:0]
    37  	return r
    38  }
    39  
    40  // Cookies returns a slice of *http.Cookie.
    41  func (r *Response) Cookies() []*http.Cookie {
    42  	return r.cookies
    43  }
    44  
    45  // SetCookie adds a Set-Cookie header to the response.
    46  func (r *Response) SetCookie(cookie *http.Cookie) {
    47  	r.cookies = append(r.cookies, cookie)
    48  	http.SetCookie(r, cookie)
    49  }
    50  
    51  func (r *Response) writeTo(w http.ResponseWriter) error {
    52  	for key, values := range r.Header() {
    53  		for _, v := range values {
    54  			w.Header().Add(key, v)
    55  		}
    56  	}
    57  	w.WriteHeader(r.resp.Code)
    58  	_, err := io.Copy(w, r.resp.Body)
    59  	responsePool.Put(r)
    60  	return err
    61  }
    62  
    63  func (r *Response) reset() {
    64  	r.StatusCode = http.StatusOK
    65  	r.resp = httptest.NewRecorder()
    66  	r.ResponseWriter = r.resp
    67  }