Repository: rollbar/rollbar-go
Pull request #67 (2019-10-03 13:57:04)

Set custom http client
Hello.

I wish there could be a way to set custom http client. Currently the only way to change http settings is either to change `http.DefaultClient` or to re-implement transport's `Send` method. Neither is handy. If you're ok with general direction will add specs and doc comments. Thanks.


Patch:
diff --git a/async_transport.go b/async_transport.go
index c80cc11..61c055a 100644
--- a/async_transport.go
+++ b/async_transport.go
@@ -7,26 +7,12 @@ import (
 // AsyncTransport is a concrete implementation of the Transport type which communicates with the
 // Rollbar API asynchronously using a buffered channel.
 type AsyncTransport struct {
-	// Rollbar access token used by this transport for communication with the Rollbar API.
-	Token string
-	// Endpoint to post items to.
-	Endpoint string
-	// Logger used to report errors when sending data to Rollbar, e.g.
-	// when the Rollbar API returns 409 Too Many Requests response.
-	// If not set, the client will use the standard log.Printf by default.
-	Logger ClientLogger
+	baseTransport
 	// Buffer is the size of the channel used for queueing asynchronous payloads for sending to
 	// Rollbar.
-	Buffer int
-	// RetryAttempts is how often to attempt to resend an item when a temporary network error occurs
-	// This defaults to DefaultRetryAttempts
-	// Set this value to 0 if you do not want retries to happen
-	RetryAttempts int
-	// PrintPayloadOnError is whether or not to output the payload to the set logger or to stderr
-	// if an error occurs during transport to the Rollbar API.
-	PrintPayloadOnError bool
-	bodyChannel         chan payload
-	waitGroup           sync.WaitGroup
+	Buffer      int
+	bodyChannel chan payload
+	waitGroup   sync.WaitGroup
 }
 
 type payload struct {
@@ -39,17 +25,19 @@ type payload struct {
 // buffer argument.
 func NewAsyncTransport(token string, endpoint string, buffer int) *AsyncTransport {
 	transport := &AsyncTransport{
-		Token:               token,
-		Endpoint:            endpoint,
-		Buffer:              buffer,
-		RetryAttempts:       DefaultRetryAttempts,
-		PrintPayloadOnError: true,
-		bodyChannel:         make(chan payload, buffer),
+		baseTransport: baseTransport{
+			Token:               token,
+			Endpoint:            endpoint,
+			RetryAttempts:       DefaultRetryAttempts,
+			PrintPayloadOnError: true,
+		},
+		bodyChannel: make(chan payload, buffer),
+		Buffer:      buffer,
 	}
 
 	go func() {
 		for p := range transport.bodyChannel {
-			err, canRetry := transport.post(p)
+			canRetry, err := transport.post(p.body)
 			if err != nil {
 				if canRetry && p.retriesLeft > 0 {
 					p.retriesLeft -= 1
@@ -111,44 +99,3 @@ func (t *AsyncTransport) Close() error {
 	t.Wait()
 	return nil
 }
-
-// SetToken updates the token to use for future API requests.
-// Any request that is currently in the queue will use this
-// updated token value. If you want to change the token without
-// affecting the items currently in the queue, use Wait first
-// to flush the queue.
-func (t *AsyncTransport) SetToken(token string) {
-	t.Token = token
-}
-
-// SetEndpoint updates the API endpoint to send items to.
-// Any request that is currently in the queue will use this
-// updated endpoint value. If you want to change the endpoint without
-// affecting the items currently in the queue, use Wait first
-// to flush the queue.
-func (t *AsyncTransport) SetEndpoint(endpoint string) {
-	t.Endpoint = endpoint
-}
-
-// SetLogger updates the logger that this transport uses for reporting errors that occur while
-// processing items.
-func (t *AsyncTransport) SetLogger(logger ClientLogger) {
-	t.Logger = logger
-}
-
-// SetRetryAttempts is how often to attempt to resend an item when a temporary network error occurs
-// This defaults to DefaultRetryAttempts
-// Set this value to 0 if you do not want retries to happen
-func (t *AsyncTransport) SetRetryAttempts(retryAttempts int) {
-	t.RetryAttempts = retryAttempts
-}
-
-// SetPrintPayloadOnError is whether or not to output the payload to stderr if an error occurs during
-// transport to the Rollbar API.
-func (t *AsyncTransport) SetPrintPayloadOnError(printPayloadOnError bool) {
-	t.PrintPayloadOnError = printPayloadOnError
-}
-
-func (t *AsyncTransport) post(p payload) (error, bool) {
-	return clientPost(t.Token, t.Endpoint, p.body, t.Logger)
-}
diff --git a/base_transport.go b/base_transport.go
new file mode 100644
index 0000000..0b16ef4
--- /dev/null
+++ b/base_transport.go
@@ -0,0 +1,108 @@
+package rollbar
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"io/ioutil"
+	"net/http"
+)
+
+type baseTransport struct {
+	// Rollbar access token used by this transport for communication with the Rollbar API.
+	Token string
+	// Endpoint to post items to.
+	Endpoint string
+	// Logger used to report errors when sending data to Rollbar, e.g.
+	// when the Rollbar API returns 409 Too Many Requests response.
+	// If not set, the client will use the standard log.Printf by default.
+	Logger ClientLogger
+	// RetryAttempts is how often to attempt to resend an item when a temporary network error occurs
+	// This defaults to DefaultRetryAttempts
+	// Set this value to 0 if you do not want retries to happen
+	RetryAttempts int
+	// PrintPayloadOnError is whether or not to output the payload to the set logger or to stderr if
+	// an error occurs during transport to the Rollbar API.
+	PrintPayloadOnError bool
+	// custom http client (http.DefaultClient used by default)
+	httpClient *http.Client
+}
+
+// SetToken updates the token to use for future API requests.
+func (t *baseTransport) SetToken(token string) {
+	t.Token = token
+}
+
+// SetEndpoint updates the API endpoint to send items to.
+func (t *baseTransport) SetEndpoint(endpoint string) {
+	t.Endpoint = endpoint
+}
+
+// SetLogger updates the logger that this transport uses for reporting errors that occur while
+// processing items.
+func (t *baseTransport) SetLogger(logger ClientLogger) {
+	t.Logger = logger
+}
+
+// SetRetryAttempts is how often to attempt to resend an item when a temporary network error occurs
+// This defaults to DefaultRetryAttempts
+// Set this value to 0 if you do not want retries to happen
+func (t *baseTransport) SetRetryAttempts(retryAttempts int) {
+	t.RetryAttempts = retryAttempts
+}
+
+// SetPrintPayloadOnError is whether or not to output the payload to stderr if an error occurs during
+// transport to the Rollbar API.
+func (t *baseTransport) SetPrintPayloadOnError(printPayloadOnError bool) {
+	t.PrintPayloadOnError = printPayloadOnError
+}
+
+// SetHTTPClient sets custom http client. http.DefaultClient is used by default
+func (t *baseTransport) SetHTTPClient(c *http.Client) {
+	t.httpClient = c
+}
+
+// getHTTPClient returns either custom client (if set) or http.DefaultClient
+func (t *baseTransport) getHTTPClient() *http.Client {
+	if t.httpClient != nil {
+		return t.httpClient
+	}
+
+	return http.DefaultClient
+}
+
+// post returns an error which indicates the type of error that occurred while attempting to
+// send the body input to the endpoint given, or nil if no error occurred. If error is not nil, the
+// boolean return parameter indicates whether the error is temporary or not. If this boolean return
+// value is true then the caller could call this function again with the same input and possibly
+// see a non-error response.
+func (t *baseTransport) post(body map[string]interface{}) (bool, error) {
+	if len(t.Token) == 0 {
+		rollbarError(t.Logger, "empty token")
+		return false, nil
+	}
+
+	jsonBody, err := json.Marshal(body)
+	if err != nil {
+		rollbarError(t.Logger, "failed to encode payload: %s", err.Error())
+		return false, err
+	}
+
+	resp, err := t.getHTTPClient().Post(t.Endpoint, "application/json", bytes.NewReader(jsonBody))
+	if err != nil {
+		rollbarError(t.Logger, "POST failed: %s", err.Error())
+		return isTemporary(err), err
+	}
+
+	io.Copy(ioutil.Discard, resp.Body)
+	resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		rollbarError(t.Logger, "received response: %s", resp.Status)
+		// http.StatusTooManyRequests is only defined in Go 1.6+ so we use 429 directly
+		isRateLimit := resp.StatusCode == 429
+		return isRateLimit, ErrHTTPError(resp.StatusCode)
+	}
+
+	return false, nil
+}
diff --git a/client.go b/client.go
index 84b4f45..c9e185a 100644
--- a/client.go
+++ b/client.go
@@ -1,9 +1,7 @@
 package rollbar
 
 import (
-	"bytes"
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -130,9 +128,9 @@ func (c *Client) SetPerson(id, username, email string) {
 	}
 
 	c.diagnostic.configuredOptions["person"] = map[string]string{
-		"Id": id,
+		"Id":       id,
 		"Username": username,
-		"Email": email,
+		"Email":    email,
 	}
 	c.configuration.person = person
 }
@@ -248,6 +246,11 @@ func (c *Client) SetPrintPayloadOnError(printPayloadOnError bool) {
 	c.Transport.SetPrintPayloadOnError(printPayloadOnError)
 }
 
+// SetHTTPClient sets custom http client. http.DefaultClient is used by default
+func (c *Client) SetHTTPClient(httpClient *http.Client) {
+	c.Transport.SetHTTPClient(httpClient)
+}
+
 // Token is the currently set Rollbar access token.
 func (c *Client) Token() string {
 	return c.configuration.token
@@ -683,51 +686,17 @@ func createConfiguration(token, environment, codeVersion, serverHost, serverRoot
 }
 
 type diagnostic struct {
-	languageVersion string
+	languageVersion   string
 	configuredOptions map[string]interface{}
 }
 
 func createDiagnostic() diagnostic {
 	return diagnostic{
-		languageVersion: runtime.Version(),
+		languageVersion:   runtime.Version(),
 		configuredOptions: map[string]interface{}{},
 	}
 }
 
-// clientPost returns an error which indicates the type of error that occured while attempting to
-// send the body input to the endpoint given, or nil if no error occurred. If error is not nil, the
-// boolean return parameter indicates whether the error is temporary or not. If this boolean return
-// value is true then the caller could call this function again with the same input and possibly
-// see a non-error response.
-func clientPost(token, endpoint string, body map[string]interface{}, logger ClientLogger) (error, bool) {
-	if len(token) == 0 {
-		rollbarError(logger, "empty token")
-		return nil, false
-	}
-
-	jsonBody, err := json.Marshal(body)
-	if err != nil {
-		rollbarError(logger, "failed to encode payload: %s", err.Error())
-		return err, false
-	}
-
-	resp, err := http.Post(endpoint, "application/json", bytes.NewReader(jsonBody))
-	if err != nil {
-		rollbarError(logger, "POST failed: %s", err.Error())
-		return err, isTemporary(err)
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != 200 {
-		rollbarError(logger, "received response: %s", resp.Status)
-		// http.StatusTooManyRequests is only defined in Go 1.6+ so we use 429 directly
-		isRateLimit := resp.StatusCode == 429
-		return ErrHTTPError(resp.StatusCode), isRateLimit
-	}
-
-	return nil, false
-}
-
 // isTemporary returns true if we should consider the error returned from http.Post to be temporary
 // in nature and possibly resolvable by simplying trying the request again.
 // https://github.com/grpc/grpc-go/blob/25b4a426b40c26c07c80af674b03db90b5bd4a60/transport/http2_client.go#L125
diff --git a/rollbar.go b/rollbar.go
index 6567792..860fc47 100644
--- a/rollbar.go
+++ b/rollbar.go
@@ -262,6 +262,11 @@ func SetPrintPayloadOnError(printPayloadOnError bool) {
 	std.SetPrintPayloadOnError(printPayloadOnError)
 }
 
+// SetHTTPClient sets custom http client. http.DefaultClient is used by default
+func SetHTTPClient(httpClient *http.Client) {
+	std.SetHTTPClient(httpClient)
+}
+
 // -- Getters
 
 // Token returns the currently set Rollbar access token on the managed Client instance.
diff --git a/sync_transport.go b/sync_transport.go
index b86fb6d..e56c2fe 100644
--- a/sync_transport.go
+++ b/sync_transport.go
@@ -3,31 +3,19 @@ package rollbar
 // SyncTransport is a concrete implementation of the Transport type which communicates with the
 // Rollbar API synchronously.
 type SyncTransport struct {
-	// Rollbar access token used by this transport for communication with the Rollbar API.
-	Token string
-	// Endpoint to post items to.
-	Endpoint string
-	// Logger used to report errors when sending data to Rollbar, e.g.
-	// when the Rollbar API returns 409 Too Many Requests response.
-	// If not set, the client will use the standard log.Printf by default.
-	Logger ClientLogger
-	// RetryAttempts is how often to attempt to resend an item when a temporary network error occurs
-	// This defaults to DefaultRetryAttempts
-	// Set this value to 0 if you do not want retries to happen
-	RetryAttempts int
-	// PrintPayloadOnError is whether or not to output the payload to the set logger or to stderr if
-	// an error occurs during transport to the Rollbar API.
-	PrintPayloadOnError bool
+	baseTransport
 }
 
 // NewSyncTransport builds a synchronous transport which sends data to the Rollbar API at the
 // specified endpoint using the given access token.
 func NewSyncTransport(token, endpoint string) *SyncTransport {
 	return &SyncTransport{
-		Token:               token,
-		Endpoint:            endpoint,
-		RetryAttempts:       DefaultRetryAttempts,
-		PrintPayloadOnError: true,
+		baseTransport{
+			Token:               token,
+			Endpoint:            endpoint,
+			RetryAttempts:       DefaultRetryAttempts,
+			PrintPayloadOnError: true,
+		},
 	}
 }
 
@@ -40,7 +28,7 @@ func (t *SyncTransport) Send(body map[string]interface{}) error {
 }
 
 func (t *SyncTransport) doSend(body map[string]interface{}, retriesLeft int) error {
-	err, canRetry := clientPost(t.Token, t.Endpoint, body, t.Logger)
+	canRetry, err := t.post(body)
 	if err != nil {
 		if !canRetry || retriesLeft <= 0 {
 			if t.PrintPayloadOnError {
@@ -60,32 +48,3 @@ func (t *SyncTransport) Wait() {}
 func (t *SyncTransport) Close() error {
 	return nil
 }
-
-// SetToken updates the token to use for future API requests.
-func (t *SyncTransport) SetToken(token string) {
-	t.Token = token
-}
-
-// SetEndpoint updates the API endpoint to send items to.
-func (t *SyncTransport) SetEndpoint(endpoint string) {
-	t.Endpoint = endpoint
-}
-
-// SetLogger updates the logger that this transport uses for reporting errors that occur while
-// processing items.
-func (t *SyncTransport) SetLogger(logger ClientLogger) {
-	t.Logger = logger
-}
-
-// SetRetryAttempts is how often to attempt to resend an item when a temporary network error occurs
-// This defaults to DefaultRetryAttempts
-// Set this value to 0 if you do not want retries to happen
-func (t *SyncTransport) SetRetryAttempts(retryAttempts int) {
-	t.RetryAttempts = retryAttempts
-}
-
-// SetPrintPayloadOnError is whether or not to output the payload to stderr if an error occurs during
-// transport to the Rollbar API.
-func (t *SyncTransport) SetPrintPayloadOnError(printPayloadOnError bool) {
-	t.PrintPayloadOnError = printPayloadOnError
-}
diff --git a/transport.go b/transport.go
index 07379a4..12faf31 100644
--- a/transport.go
+++ b/transport.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"net/http"
 	"os"
 )
 
@@ -37,6 +38,8 @@ type Transport interface {
 	SetRetryAttempts(retryAttempts int)
 	// Set whether to print the payload to the set logger or to stderr upon failing to send.
 	SetPrintPayloadOnError(printPayloadOnError bool)
+	// Sets custom http client. http.DefaultClient is used by default
+	SetHTTPClient(httpClient *http.Client)
 }
 
 // ClientLogger is the interface used by the rollbar Client/Transport to report problems.


Test patch:
diff --git a/client_test.go b/client_test.go
index 7eaf3eb..13e5592 100644
--- a/client_test.go
+++ b/client_test.go
@@ -1,14 +1,15 @@
 package rollbar_test
 
 import (
+	"context"
 	"errors"
+	"fmt"
 	"github.com/rollbar/rollbar-go"
-	"regexp"
-	"context"
+	"net/http"
 	"reflect"
-	"testing"
+	"regexp"
 	"strings"
-	"fmt"
+	"testing"
 )
 
 type TestTransport struct {
@@ -29,6 +30,7 @@ func (t *TestTransport) SetEndpoint(_e string)             {}
 func (t *TestTransport) SetLogger(_l rollbar.ClientLogger) {}
 func (t *TestTransport) SetRetryAttempts(_r int)           {}
 func (t *TestTransport) SetPrintPayloadOnError(_p bool)    {}
+func (t *TestTransport) SetHTTPClient(_c *http.Client)     {}
 func (t *TestTransport) Send(body map[string]interface{}) error {
 	t.Body = body
 	return nil
diff --git a/rollbar_test.go b/rollbar_test.go
index 0892094..df58dc1 100644
--- a/rollbar_test.go
+++ b/rollbar_test.go
@@ -4,9 +4,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"net/http"
 	"os"
 	"runtime"
+	"strings"
 	"testing"
 )
 
@@ -629,3 +631,47 @@ func TestSetStackTracer(t *testing.T) {
 		t.Errorf("custom tracer got the wrong trace")
 	}
 }
+
+type roundTripFunc func(r *http.Request) (*http.Response, error)
+
+func (s roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
+	return s(r)
+}
+
+func TestSetHttpClient(t *testing.T) {
+	used := false
+	c := &http.Client{
+		Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+			used = true
+			return &http.Response{
+				StatusCode: http.StatusOK,
+				Body:       ioutil.NopCloser(strings.NewReader("")),
+			}, nil
+		}),
+	}
+
+	client := NewAsync("example", "test", "0.0.0", "", "")
+	client.SetHTTPClient(c)
+
+	err := client.Transport.Send(map[string]interface{}{})
+	client.Wait()
+	if err != nil {
+		t.Fatal("failed to send body:", err.Error())
+	}
+
+	if !used {
+		t.Fatal("custom http client had not been invoked")
+	}
+
+	used = false
+	client = NewSync("example", "test", "0.0.0", "", "")
+	client.SetHTTPClient(c)
+
+	if err := client.Transport.Send(map[string]interface{}{}); err != nil {
+		t.Fatal("failed to send body:", err.Error())
+	}
+
+	if !used {
+		t.Fatal("custom http client had not been invoked")
+	}
+}