From b538053e1bfd5e975ff1e80ddb4684feb871b7f9 Mon Sep 17 00:00:00 2001
From: reecerussell <me@reece-russell.co.uk>
Date: Mon, 24 Mar 2025 15:48:33 +0000
Subject: [PATCH] feat: added the last response as argument to PrepareRetry

---
 client.go      |  7 ++++---
 client_test.go | 20 +++++++++++++++++++-
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/client.go b/client.go
index efee53c..70c6c7d 100644
--- a/client.go
+++ b/client.go
@@ -398,8 +398,9 @@ type Backoff func(min, max time.Duration, attemptNum int, resp *http.Response) t
 // attempted. If overriding this, be sure to close the body if needed.
 type ErrorHandler func(resp *http.Response, err error, numTries int) (*http.Response, error)
 
-// PrepareRetry is called before retry operation. It can be used for example to re-sign the request
-type PrepareRetry func(req *http.Request) error
+// PrepareRetry is called before retry operation. It can be used for example to re-sign the request or
+// update it depending on the previous failed response.
+type PrepareRetry func(req *http.Request, resp *http.Response) error
 
 // Client is used to make HTTP requests. It adds additional functionality
 // like automatic retries to tolerate minor outages.
@@ -779,7 +780,7 @@ func (c *Client) Do(req *Request) (*http.Response, error) {
 		req.Request = &httpreq
 
 		if c.PrepareRetry != nil {
-			if err := c.PrepareRetry(req.Request); err != nil {
+			if err := c.PrepareRetry(req.Request, resp); err != nil {
 				prepareErr = err
 				break
 			}
diff --git a/client_test.go b/client_test.go
index cc05e91..89db1c1 100644
--- a/client_test.go
+++ b/client_test.go
@@ -369,9 +369,14 @@ func TestClient_Do_WithPrepareRetry(t *testing.T) {
 	}
 
 	var prepareChecks int
-	client.PrepareRetry = func(req *http.Request) error {
+	client.PrepareRetry = func(req *http.Request, resp *http.Response) error {
 		prepareChecks++
 		req.Header.Set("foo", strconv.Itoa(prepareChecks))
+
+		lastCount := resp.Header.Get("bar")
+		if lastCount != "" && lastCount != strconv.Itoa(prepareChecks) {
+			return errors.New("prepare check failed")
+		}
 		return nil
 	}
 
@@ -431,6 +436,19 @@ func TestClient_Do_WithPrepareRetry(t *testing.T) {
 			expectedChecks:        4,
 			expectedPrepareChecks: 1,
 		},
+		{
+			name: "request succeeds when matches with response",
+			handler: func(resp *http.Response) error {
+				if shouldSucceed {
+					resp.Header.Set("bar", strconv.Itoa(prepareChecks))
+					return nil
+				}
+				shouldSucceed = true
+				return errors.New("retryable failure")
+			},
+			expectedChecks:        4,
+			expectedPrepareChecks: 1,
+		},
 	}
 
 	for _, tt := range tests {