From f0c13f63f9fe14c0d9ff64218ce0968f25a124f2 Mon Sep 17 00:00:00 2001
From: Andrew Starr-Bochicchio <andrewsomething@users.noreply.github.com>
Date: Mon, 14 Aug 2023 15:55:48 -0400
Subject: [PATCH] Allow configuring go-retryablehttp.Logger (#626)

---
 godo.go      |  9 +++++++--
 godo_test.go | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/godo.go b/godo.go
index d5e90217..cc8e89c2 100644
--- a/godo.go
+++ b/godo.go
@@ -114,8 +114,9 @@ type Client struct {
 // Only the custom HTTP client's custom transport and timeout will be maintained.
 type RetryConfig struct {
 	RetryMax     int
-	RetryWaitMin *float64 // Minimum time to wait
-	RetryWaitMax *float64 // Maximum time to wait
+	RetryWaitMin *float64    // Minimum time to wait
+	RetryWaitMax *float64    // Maximum time to wait
+	Logger       interface{} // Customer logger instance. Must implement either go-retryablehttp.Logger or go-retryablehttp.LeveledLogger
 }
 
 // RequestCompletionCallback defines the type of the request callback function
@@ -307,6 +308,9 @@ func New(httpClient *http.Client, opts ...ClientOpt) (*Client, error) {
 			retryableClient.RetryWaitMax = time.Duration(*c.RetryConfig.RetryWaitMax * float64(time.Second))
 		}
 
+		// By default this is nil and does not log.
+		retryableClient.Logger = c.RetryConfig.Logger
+
 		// if timeout is set, it is maintained before overwriting client with StandardClient()
 		retryableClient.HTTPClient.Timeout = c.HTTPClient.Timeout
 
@@ -373,6 +377,7 @@ func WithRetryAndBackoffs(retryConfig RetryConfig) ClientOpt {
 		c.RetryConfig.RetryMax = retryConfig.RetryMax
 		c.RetryConfig.RetryWaitMax = retryConfig.RetryWaitMax
 		c.RetryConfig.RetryWaitMin = retryConfig.RetryWaitMin
+		c.RetryConfig.Logger = retryConfig.Logger
 		return nil
 	}
 }
diff --git a/godo_test.go b/godo_test.go
index 0badba83..323eba72 100644
--- a/godo_test.go
+++ b/godo_test.go
@@ -1,9 +1,11 @@
 package godo
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io/ioutil"
+	"log"
 	"net/http"
 	"net/http/httptest"
 	"net/http/httputil"
@@ -652,6 +654,53 @@ func TestWithRetryAndBackoffs(t *testing.T) {
 
 }
 
+func TestWithRetryAndBackoffsLogger(t *testing.T) {
+	// Mock server which always responds 500.
+	setup()
+	defer teardown()
+
+	url, _ := url.Parse(server.URL)
+	mux.HandleFunc("/foo", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusAccepted)
+	})
+
+	tokenSrc := oauth2.StaticTokenSource(&oauth2.Token{
+		AccessToken: "new_token",
+	})
+
+	oauth_client := oauth2.NewClient(oauth2.NoContext, tokenSrc)
+
+	var buf bytes.Buffer
+	retryConfig := RetryConfig{
+		RetryMax: 3,
+		Logger:   log.New(&buf, "", 0),
+	}
+
+	// Create the client. Use short retry windows so we fail faster.
+	client, err := New(oauth_client, WithRetryAndBackoffs(retryConfig))
+	client.BaseURL = url
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Create the request
+	req, err := client.NewRequest(ctx, http.MethodGet, "/foo", nil)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	_, err = client.Do(context.Background(), req, nil)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	got := buf.String()
+	expected := fmt.Sprintf("[DEBUG] GET %s/foo\n", url)
+	if expected != got {
+		t.Fatalf("expected: %s; got: %s", expected, got)
+	}
+}
+
 func checkCurrentPage(t *testing.T, resp *Response, expectedPage int) {
 	links := resp.Links
 	p, err := links.CurrentPage()