Allow users to specify their own client implementation used by the library (#73)

* Allow users to specify their own client implementation used by the library

* fix typos

* add tests

Author: nikola-jokic

common_client.go

@@ -75,12 +75,20 @@ func sendRequest(c *http.Client, req *http.Request) (*http.Response, error) {
 }
 
 type httpClientOption struct {
-	logger                *slog.Logger
-	retryMax              int
-	retryWaitMax          time.Duration
+	logger *slog.Logger
+
+	// Options for built-in retryable HTTP client.
+	// Ignored if a custom retryable HTTP client is provided via WithRetryableHTTPClint.
+	retryMax     int
+	retryWaitMax time.Duration
+
+	// fields added to the transport if specified
 	rootCAs               *x509.CertPool
 	tlsInsecureSkipVerify bool
 	proxyFunc             ProxyFunc
+	timeout               time.Duration
+
+	retryableHTTPClient *retryablehttp.Client
 }
 
 func (o *httpClientOption) defaults() {
@@ -93,14 +101,28 @@ func (o *httpClientOption) defaults() {
 	if o.retryWaitMax == 0 {
 		o.retryWaitMax = 30 * time.Second
 	}
+	if o.timeout == 0 {
+		o.timeout = 5 * time.Minute
+	}
 }
 
 func (o *httpClientOption) newRetryableHTTPClient() (*retryablehttp.Client, error) {
-	retryClient := retryablehttp.NewClient()
-	retryClient.Logger = o.logger
-	retryClient.RetryMax = o.retryMax
-	retryClient.RetryWaitMax = o.retryWaitMax
-	retryClient.HTTPClient.Timeout = 5 * time.Minute // timeout must be > 1m to accomodate long polling
+	var retryClient *retryablehttp.Client
+	if o.retryableHTTPClient != nil {
+		retryClient = o.retryableHTTPClient
+	} else {
+		retryClient = retryablehttp.NewClient()
+		retryClient.RetryMax = o.retryMax
+		retryClient.RetryWaitMax = o.retryWaitMax
+	}
+
+	if retryClient.HTTPClient.Timeout == 0 {
+		retryClient.HTTPClient.Timeout = o.timeout
+	}
+
+	if retryClient.Logger == nil {
+		retryClient.Logger = o.logger
+	}
 
 	transport, ok := retryClient.HTTPClient.Transport.(*http.Transport)
 	if !ok {
@@ -120,7 +142,9 @@ func (o *httpClientOption) newRetryableHTTPClient() (*retryablehttp.Client, erro
 		transport.TLSClientConfig.InsecureSkipVerify = true
 	}
 
-	transport.Proxy = o.proxyFunc
+	if o.proxyFunc != nil {
+		transport.Proxy = o.proxyFunc
+	}
 
 	retryClient.HTTPClient.Transport = transport
 
@@ -145,6 +169,14 @@ func (c *commonClient) setUserAgent() {
 // HTTPOption defines a functional option for configuring the Client.
 type HTTPOption func(*httpClientOption)
 
+// WithRetryableHTTPClint allows users to provide a custom retryable HTTP client.
+// If not set, a default client will be used with the specified retry and timeout settings.
+func WithRetryableHTTPClint(client *retryablehttp.Client) HTTPOption {
+	return func(c *httpClientOption) {
+		c.retryableHTTPClient = client
+	}
+}
+
 // WithLogger sets a custom logger for the Client.
 // If nil is passed, a discard logger will be used.
 func WithLogger(logger *slog.Logger) HTTPOption {
@@ -190,3 +222,10 @@ func WithProxy(proxyFunc ProxyFunc) HTTPOption {
 		c.proxyFunc = proxyFunc
 	}
 }
+
+// WithTimeout sets a timeout for the Client.
+func WithTimeout(duration time.Duration) HTTPOption {
+	return func(c *httpClientOption) {
+		c.timeout = duration
+	}
+}

common_client_test.go

@@ -7,8 +7,10 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"testing"
+	"time"
 
 	"github.com/actions/scaleset/internal/testserver"
+	"github.com/hashicorp/go-retryablehttp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/http/httpproxy"
@@ -153,3 +155,88 @@ func TestUserAgent(t *testing.T) {
 
 	assert.Equal(t, want, got)
 }
+
+// TestWithRetryableHTTPClient verifies that a custom retryable HTTP client
+// provided via WithRetryableHTTPClient is actually used instead of the built-in one
+func TestWithRetryableHTTPClient(t *testing.T) {
+	t.Run("uses custom retryable client instead of built-in", func(t *testing.T) {
+		attemptCount := 0
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			attemptCount++
+			if attemptCount == 1 {
+				w.WriteHeader(http.StatusServiceUnavailable)
+				return
+			}
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte(`{"result": "success"}`))
+		}))
+		defer server.Close()
+
+		// Create a custom retryable HTTP client with specific retry configuration
+		customRetryClient := retryablehttp.NewClient()
+		customRetryClient.RetryMax = 3
+		customRetryClient.RetryWaitMax = 10 * time.Millisecond
+
+		// Create options with the custom retryable client
+		opts := defaultHTTPClientOption()
+		WithRetryableHTTPClint(customRetryClient)(&opts)
+
+		// Verify that the custom client is set in options
+		assert.NotNil(t, opts.retryableHTTPClient)
+		assert.Equal(t, customRetryClient, opts.retryableHTTPClient)
+
+		// Create the common client with custom retryable client
+		client := newCommonClient(testSystemInfo, opts)
+
+		// Make a request that will trigger a retry
+		req, err := http.NewRequest("GET", server.URL, nil)
+		require.NoError(t, err)
+
+		resp, err := client.do(req)
+		require.NoError(t, err)
+
+		// Should succeed after retry
+		assert.Equal(t, http.StatusOK, resp.StatusCode)
+		assert.Equal(t, 2, attemptCount)
+
+		// Verify that the client used is the custom one by checking newRetryableHTTPClient
+		retrievedRetryClient, err := client.newRetryableHTTPClient()
+		require.NoError(t, err)
+		assert.Equal(t, customRetryClient, retrievedRetryClient, "should return the custom retryable client")
+	})
+
+	t.Run("respects custom client's retry configuration over built-in defaults", func(t *testing.T) {
+		attemptCount := 0
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			attemptCount++
+			w.WriteHeader(http.StatusServiceUnavailable)
+		}))
+		defer server.Close()
+
+		// Create custom client with limited retries
+		customRetryClient := retryablehttp.NewClient()
+		customRetryClient.RetryMax = 1 // Only 1 retry (2 total attempts)
+		customRetryClient.RetryWaitMax = 5 * time.Millisecond
+
+		opts := defaultHTTPClientOption()
+		WithRetryableHTTPClint(customRetryClient)(&opts)
+
+		client := newCommonClient(testSystemInfo, opts)
+
+		req, err := http.NewRequest("GET", server.URL, nil)
+		require.NoError(t, err)
+
+		resp, err := client.do(req)
+		// When all retries are exhausted with a retryable error, the client gives up
+		// and an error is returned
+		if err != nil {
+			// Expected: request failed after exhausting retries
+			assert.Contains(t, err.Error(), "giving up after 2 attempt(s)")
+		} else {
+			// Or the final response is returned
+			assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode)
+		}
+		// Should have tried 1 initial + 1 retry = 2 times total
+		assert.Equal(t, 2, attemptCount)
+	})
+}