Allow injection of HTTP transport to enable HTTP replayer pattern

client/client.go

@@ -35,30 +35,34 @@ func New(cfg *config.Config) (*DatabricksClient, error) {
 	httpTimeout := time.Duration(orDefault(cfg.HTTPTimeoutSeconds, 60)) * time.Second
 	rateLimiter := rate.NewLimiter(rate.Limit(orDefault(cfg.RateLimitPerSecond, 15)), 1)
 	debugTruncateBytes := orDefault(cfg.DebugTruncateBytes, 96)
+	httpTransport := cfg.HTTPTransport
+	if httpTransport == nil {
+		httpTransport = &http.Transport{
+			Proxy: http.ProxyFromEnvironment,
+			DialContext: (&net.Dialer{
+				Timeout:   30 * time.Second,
+				KeepAlive: 30 * time.Second,
+			}).DialContext,
+			ForceAttemptHTTP2:     true,
+			MaxIdleConns:          100,
+			MaxIdleConnsPerHost:   runtime.GOMAXPROCS(0) + 1,
+			IdleConnTimeout:       180 * time.Second,
+			TLSHandshakeTimeout:   30 * time.Second,
+			ExpectContinueTimeout: 1 * time.Second,
+			TLSClientConfig: &tls.Config{
+				InsecureSkipVerify: cfg.InsecureSkipVerify,
+			},
+		}
+	}
 	return &DatabricksClient{
 		Config:             cfg,
 		debugHeaders:       cfg.DebugHeaders,
 		debugTruncateBytes: debugTruncateBytes,
 		retryTimeout:       retryTimeout,
 		rateLimiter:        rateLimiter,
 		httpClient: &http.Client{
-			Timeout: httpTimeout,
-			Transport: &http.Transport{
-				Proxy: http.ProxyFromEnvironment,
-				DialContext: (&net.Dialer{
-					Timeout:   30 * time.Second,
-					KeepAlive: 30 * time.Second,
-				}).DialContext,
-				ForceAttemptHTTP2:     true,
-				MaxIdleConns:          100,
-				MaxIdleConnsPerHost:   runtime.GOMAXPROCS(0) + 1,
-				IdleConnTimeout:       180 * time.Second,
-				TLSHandshakeTimeout:   30 * time.Second,
-				ExpectContinueTimeout: 1 * time.Second,
-				TLSClientConfig: &tls.Config{
-					InsecureSkipVerify: cfg.InsecureSkipVerify,
-				},
-			},
+			Timeout:   httpTimeout,
+			Transport: httpTransport,
 		},
 	}, nil
 }

client/client_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/logger"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"golang.org/x/time/rate"
 )
 
@@ -727,3 +728,23 @@ func TestRetryGetRequest(t *testing.T) {
 	assert.Equal(t, "succeeded", respBytes.String())
 	assert.True(t, succeed)
 }
+
+func (cb hc) RoundTrip(r *http.Request) (*http.Response, error) {
+	return cb(r)
+}
+
+func TestHttpTransport(t *testing.T) {
+	calledMock := false
+	cfg := config.NewMockConfig(func(r *http.Request) error { return nil })
+	cfg.HTTPTransport = hc(func(r *http.Request) (*http.Response, error) {
+		calledMock = true
+		return &http.Response{Request: r}, nil
+	})
+	client, err := New(cfg)
+	require.NoError(t, err)
+
+	err = client.Do(context.Background(), "GET", "/a", nil, nil, bytes.Buffer{})
+	require.NoError(t, err)
+
+	assert.True(t, calledMock)
+}

config/config.go

@@ -105,6 +105,9 @@ type Config struct {
 	// Number of seconds to keep retrying HTTP requests. Default is 300 (5 minutes)
 	RetryTimeoutSeconds int `name:"retry_timeout_seconds" auth:"-"`
 
+	// HTTPTransport can be overriden for unit testing and together with tooling like https://github.com/google/go-replayers
+	HTTPTransport http.RoundTripper
+
 	Loaders []Loader
 
 	// marker for configuration resolving