github.com/blend/go-sdk@v1.20220411.3/r2/r2test/opt_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 r2test
     9  
    10  import (
    11  	"fmt"
    12  	"net/http"
    13  	"net/http/httptest"
    14  	"net/url"
    15  
    16  	"github.com/blend/go-sdk/r2"
    17  )
    18  
    19  // OptMockResponseString mocks a string response.
    20  func OptMockResponseString(response string) r2.Option {
    21  	return OptMockResponse(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
    22  		rw.WriteHeader(http.StatusOK)
    23  		fmt.Fprint(rw, response)
    24  	}))
    25  }
    26  
    27  // OptMockResponseStringStatus mocks a string response with a given status code.
    28  func OptMockResponseStringStatus(statusCode int, response string) r2.Option {
    29  	return OptMockResponse(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
    30  		rw.WriteHeader(statusCode)
    31  		fmt.Fprint(rw, response)
    32  	}))
    33  }
    34  
    35  // OptMockResponse mocks a response by creating an httptest server.
    36  func OptMockResponse(handler http.Handler) r2.Option {
    37  	return func(r *r2.Request) error {
    38  		server := httptest.NewServer(handler)
    39  		parsedURL, _ := url.Parse(server.URL)
    40  		if r.Request.URL == nil {
    41  			// unclear if this is even possible
    42  			r.Request.URL = parsedURL
    43  		} else {
    44  			r.Request.URL.Scheme = parsedURL.Scheme
    45  			r.Request.URL.Host = parsedURL.Host
    46  		}
    47  
    48  		if r.Closer != nil {
    49  			originalCloser := r.Closer
    50  			r.Closer = func() error {
    51  				server.Close()
    52  				return originalCloser()
    53  			}
    54  		} else {
    55  			r.Closer = func() error {
    56  				server.Close()
    57  				return nil
    58  			}
    59  		}
    60  		return nil
    61  	}
    62  }