github.com/bigcommerce/nomad@v0.9.3-bc/testutil/responsewriter.go (about)

     1  package testutil
     2  
     3  import (
     4  	"net/http"
     5  	"net/http/httptest"
     6  	"sync"
     7  )
     8  
     9  // assert ResponseRecorder implements the http.ResponseWriter interface
    10  var _ http.ResponseWriter = (*ResponseRecorder)(nil)
    11  
    12  // ResponseRecorder implements a ResponseWriter which can be written to and
    13  // read from concurrently. For use in testing streaming APIs where
    14  // httptest.ResponseRecorder is unsafe for concurrent access. Uses
    15  // httptest.ResponseRecorder internally and exposes most of the functionality.
    16  type ResponseRecorder struct {
    17  	rr *httptest.ResponseRecorder
    18  	mu sync.Mutex
    19  }
    20  
    21  func NewResponseRecorder() *ResponseRecorder {
    22  	return &ResponseRecorder{
    23  		rr: httptest.NewRecorder(),
    24  	}
    25  }
    26  
    27  // Flush sets Flushed=true.
    28  func (r *ResponseRecorder) Flush() {
    29  	r.mu.Lock()
    30  	defer r.mu.Unlock()
    31  	r.rr.Flush()
    32  }
    33  
    34  // Flushed returns true if Flush has been called.
    35  func (r *ResponseRecorder) Flushed() bool {
    36  	r.mu.Lock()
    37  	defer r.mu.Unlock()
    38  	return r.rr.Flushed
    39  }
    40  
    41  // Header returns the response headers. Readers should call HeaderMap() to
    42  // avoid races due to the server concurrently mutating headers.
    43  func (r *ResponseRecorder) Header() http.Header {
    44  	r.mu.Lock()
    45  	defer r.mu.Unlock()
    46  	return r.rr.Header()
    47  }
    48  
    49  // HeaderMap returns the HTTP headers written before WriteHeader was called.
    50  func (r *ResponseRecorder) HeaderMap() http.Header {
    51  	r.mu.Lock()
    52  	defer r.mu.Unlock()
    53  	return r.rr.HeaderMap
    54  }
    55  
    56  // Write to the underlying response buffer. Safe to call concurrent with Read.
    57  func (r *ResponseRecorder) Write(p []byte) (int, error) {
    58  	r.mu.Lock()
    59  	defer r.mu.Unlock()
    60  	return r.rr.Body.Write(p)
    61  }
    62  
    63  // WriteHeader sets the response code and freezes the headers returned by
    64  // HeaderMap. Safe to call concurrent with Read and HeaderMap.
    65  func (r *ResponseRecorder) WriteHeader(statusCode int) {
    66  	r.mu.Lock()
    67  	defer r.mu.Unlock()
    68  	r.rr.WriteHeader(statusCode)
    69  }
    70  
    71  // Read available response bytes. Safe to call concurrently with Write().
    72  func (r *ResponseRecorder) Read(p []byte) (int, error) {
    73  	r.mu.Lock()
    74  	defer r.mu.Unlock()
    75  	return r.rr.Body.Read(p)
    76  }