[Go Client Mock]

tclient.go

package tclient

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
)

var unexpectRoundTrip = func(req *http.Request) *http.Response {
	return &http.Response{
		StatusCode: 500,
		Body:       ioutil.NopCloser(bytes.NewBufferString(fmt.Sprintf("unexpect request, %+v", req))),
		Header:     make(http.Header),
	}
}

type MockRoundTrips map[string]func(req *http.Request) *http.Response

func (f MockRoundTrips) RoundTrip(req *http.Request) (*http.Response, error) {
	rtf, exists := f[req.URL.String()]
	if !exists {
		return unexpectRoundTrip(req), nil
	}

	return rtf(req), nil
}

func NewMock() MockRoundTrips {
	return MockRoundTrips{}
}

func (m MockRoundTrips) Mock(url, body string, statusCode int) MockRoundTrips {
	m[url] = func(req *http.Request) *http.Response {
		return &http.Response{
			StatusCode: statusCode,
			// Send response to be tested
			Body: ioutil.NopCloser(bytes.NewBufferString(body)),
			// Must be set to non-nil value or it panics
			Header: make(http.Header),
		}
	}

	return m
}

func (m MockRoundTrips) Client() *http.Client {
	return &http.Client{
		Transport: m,
	}
}