Exposes configuration function to be able to set the Transport - Added test

Signed-off-by: Leandro Deveikis <[email protected]>

Author: leandro-deveikis

connector.go

@@ -214,3 +214,10 @@ func WithAuthenticator(authr auth.Authenticator) connOption {
 		c.Authenticator = authr
 	}
 }
+
+// WithTransport sets up the transport configuration to be used by the httpclient.
+func WithTransport(t http.RoundTripper) connOption {
+	return func(c *config.Config) {
+		c.Transport = t
+	}
+}

connector_test.go

@@ -1,6 +1,7 @@
 package dbsql
 
 import (
+	"net/http"
 	"testing"
 	"time"
 
@@ -22,6 +23,7 @@ func TestNewConnector(t *testing.T) {
 		schema := "schema-string"
 		userAgentEntry := "user-agent"
 		sessionParams := map[string]string{"key": "value"}
+		roundTripper := mockRoundTripper{}
 		con, err := NewConnector(
 			WithServerHostname(host),
 			WithPort(port),
@@ -33,6 +35,7 @@ func TestNewConnector(t *testing.T) {
 			WithUserAgentEntry(userAgentEntry),
 			WithSessionParams(sessionParams),
 			WithRetries(10, 3*time.Second, 60*time.Second),
+			WithTransport(roundTripper),
 		)
 		expectedUserConfig := config.UserConfig{
 			Host:           host,
@@ -50,6 +53,7 @@ func TestNewConnector(t *testing.T) {
 			RetryMax:       10,
 			RetryWaitMin:   3 * time.Second,
 			RetryWaitMax:   60 * time.Second,
+			Transport:      roundTripper,
 		}
 		expectedCfg := config.WithDefaults()
 		expectedCfg.DriverVersion = DriverVersion
@@ -127,3 +131,11 @@ func TestNewConnector(t *testing.T) {
 		assert.Equal(t, expectedCfg, coni.cfg)
 	})
 }
+
+type mockRoundTripper struct{}
+
+var _ http.RoundTripper = mockRoundTripper{}
+
+func (m mockRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
+	return &http.Response{StatusCode: 200}, nil
+}

internal/client/client.go

@@ -319,7 +319,7 @@ func isRetryableServerResponse(resp *http.Response) bool {
 }
 
 type Transport struct {
-	Base  *http.Transport
+	Base  http.RoundTripper
 	Authr auth.Authenticator
 	trace bool
 }
@@ -396,10 +396,20 @@ func PooledClient(cfg *config.Config) *http.Client {
 	if cfg.Authenticator == nil {
 		return nil
 	}
-	tr := &Transport{
-		Base:  PooledTransport(),
-		Authr: cfg.Authenticator,
+
+	var tr *Transport
+	if cfg.Transport != nil {
+		tr = &Transport{
+			Base:  cfg.Transport,
+			Authr: cfg.Authenticator,
+		}
+	} else {
+		tr = &Transport{
+			Base:  PooledTransport(),
+			Authr: cfg.Authenticator,
+		}
 	}
+
 	return &http.Client{
 		Transport: tr,
 		Timeout:   cfg.ClientTimeout,

internal/config/config.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"net/http"
 	"net/url"
 	"strconv"
 	"strings"
@@ -97,6 +98,7 @@ type UserConfig struct {
 	RetryWaitMin   time.Duration
 	RetryWaitMax   time.Duration
 	RetryMax       int
+	Transport      http.RoundTripper
 }
 
 // DeepCopy returns a true deep copy of UserConfig
@@ -135,6 +137,7 @@ func (ucfg UserConfig) DeepCopy() UserConfig {
 		RetryWaitMin:   ucfg.RetryWaitMin,
 		RetryWaitMax:   ucfg.RetryWaitMax,
 		RetryMax:       ucfg.RetryMax,
+		Transport:      ucfg.Transport,
 	}
 }