|
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566 |
- package martiantest
-
- import (
- "net/http"
-
- "github.com/google/martian/proxyutil"
- )
-
- // Transport is an http.RoundTripper for testing.
- type Transport struct {
- rtfunc func(*http.Request) (*http.Response, error)
- }
-
- // NewTransport builds a new transport that will respond with a 200 OK
- // response.
- func NewTransport() *Transport {
- tr := &Transport{}
- tr.Respond(200)
-
- return tr
- }
-
- // Respond sets the transport to respond with response with statusCode.
- func (tr *Transport) Respond(statusCode int) {
- tr.rtfunc = func(req *http.Request) (*http.Response, error) {
- // Force CONNECT requests to 200 to test CONNECT with downstream proxy.
- if req.Method == "CONNECT" {
- statusCode = 200
- }
-
- res := proxyutil.NewResponse(statusCode, nil, req)
-
- return res, nil
- }
- }
-
- // RespondError sets the transport to respond with an error on round trip.
- func (tr *Transport) RespondError(err error) {
- tr.rtfunc = func(*http.Request) (*http.Response, error) {
- return nil, err
- }
- }
-
- // CopyHeaders sets the transport to respond with a 200 OK response with
- // headers copied from the request to the response verbatim.
- func (tr *Transport) CopyHeaders(names ...string) {
- tr.rtfunc = func(req *http.Request) (*http.Response, error) {
- res := proxyutil.NewResponse(200, nil, req)
-
- for _, n := range names {
- res.Header.Set(n, req.Header.Get(n))
- }
-
- return res, nil
- }
- }
-
- // Func sets the transport to use the rtfunc.
- func (tr *Transport) Func(rtfunc func(*http.Request) (*http.Response, error)) {
- tr.rtfunc = rtfunc
- }
-
- // RoundTrip runs the stored round trip func and returns the response.
- func (tr *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
- return tr.rtfunc(req)
- }
|