github.com/blend/go-sdk@v1.20220411.3/webutil/mock_response.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package webutil
     9  
    10  import (
    11  	"bytes"
    12  	"io"
    13  	"net/http"
    14  )
    15  
    16  var (
    17  	_ http.ResponseWriter = (*MockResponseWriter)(nil)
    18  	_ http.Flusher        = (*MockResponseWriter)(nil)
    19  )
    20  
    21  // NewMockResponse returns a mocked response writer.
    22  func NewMockResponse(buffer io.Writer) *MockResponseWriter {
    23  	return &MockResponseWriter{
    24  		innerWriter: buffer,
    25  		contents:    new(bytes.Buffer),
    26  		headers:     http.Header{},
    27  	}
    28  }
    29  
    30  // MockResponseWriter is an object that satisfies response writer but uses an internal buffer.
    31  type MockResponseWriter struct {
    32  	innerWriter   io.Writer
    33  	contents      *bytes.Buffer
    34  	statusCode    int
    35  	contentLength int
    36  	headers       http.Header
    37  }
    38  
    39  // Write writes data and adds to ContentLength.
    40  func (res *MockResponseWriter) Write(buffer []byte) (int, error) {
    41  	bytesWritten, err := res.innerWriter.Write(buffer)
    42  	res.contentLength += bytesWritten
    43  	defer func() {
    44  		res.contents.Write(buffer)
    45  	}()
    46  	return bytesWritten, err
    47  }
    48  
    49  // Header returns the response headers.
    50  func (res *MockResponseWriter) Header() http.Header {
    51  	return res.headers
    52  }
    53  
    54  // WriteHeader sets the status code.
    55  func (res *MockResponseWriter) WriteHeader(statusCode int) {
    56  	res.statusCode = statusCode
    57  }
    58  
    59  // InnerResponse returns the backing httpresponse writer.
    60  func (res *MockResponseWriter) InnerResponse() http.ResponseWriter {
    61  	return res
    62  }
    63  
    64  // StatusCode returns the status code.
    65  func (res *MockResponseWriter) StatusCode() int {
    66  	return res.statusCode
    67  }
    68  
    69  // ContentLength returns the content length.
    70  func (res *MockResponseWriter) ContentLength() int {
    71  	return res.contentLength
    72  }
    73  
    74  // Bytes returns the raw response.
    75  func (res *MockResponseWriter) Bytes() []byte {
    76  	return res.contents.Bytes()
    77  }
    78  
    79  // Flush is a no-op.
    80  func (res *MockResponseWriter) Flush() {}
    81  
    82  // Close is a no-op.
    83  func (res *MockResponseWriter) Close() error {
    84  	return nil
    85  }