github.com/hashicorp/nomad/api@v0.0.0-20240306165712-3193ac204f65/internal/testutil/responsewriter.go (about)

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