From 207df31d0280c8ede938a3eba257bbf8539758c4 Mon Sep 17 00:00:00 2001
From: Chris Marchesi <chrism@vancluevertech.com>
Date: Fri, 23 Apr 2021 17:55:05 -0700
Subject: [PATCH] add StatusServiceUnavailable to Retry-After check in
 DefaultBackoff

This adds StatusServiceUnavailable (error 503) to the codes checked for
the Retry-After header in DefaultBackoff.

This is a valid header to honor for 503 responses. See:
  https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After
---
 client.go      |  2 +-
 client_test.go | 50 +++++++++++++++++++++++++++-----------------------
 2 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/client.go b/client.go
index 79dc931..adbdd92 100644
--- a/client.go
+++ b/client.go
@@ -471,7 +471,7 @@ func baseRetryPolicy(resp *http.Response, err error) (bool, error) {
 // seconds the server states it may be ready to process more requests from this client.
 func DefaultBackoff(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
 	if resp != nil {
-		if resp.StatusCode == http.StatusTooManyRequests {
+		if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable {
 			if s, ok := resp.Header["Retry-After"]; ok {
 				if sleep, err := strconv.ParseInt(s[0], 10, 64); err == nil {
 					return time.Second * time.Duration(sleep)
diff --git a/client_test.go b/client_test.go
index b8a3d9d..082b407 100644
--- a/client_test.go
+++ b/client_test.go
@@ -523,35 +523,39 @@ func TestClient_CheckRetry(t *testing.T) {
 	}
 }
 
-func TestClient_DefaultBackoff429TooManyRequest(t *testing.T) {
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Retry-After", "2")
-		http.Error(w, "test_429_body", http.StatusTooManyRequests)
-	}))
-	defer ts.Close()
+func TestClient_DefaultBackoff(t *testing.T) {
+	for _, code := range []int{http.StatusTooManyRequests, http.StatusServiceUnavailable} {
+		t.Run(fmt.Sprintf("http_%d", code), func(t *testing.T) {
+			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Retry-After", "2")
+				http.Error(w, fmt.Sprintf("test_%d_body", code), code)
+			}))
+			defer ts.Close()
 
-	client := NewClient()
+			client := NewClient()
 
-	var retryAfter time.Duration
-	retryable := false
+			var retryAfter time.Duration
+			retryable := false
 
-	client.CheckRetry = func(_ context.Context, resp *http.Response, err error) (bool, error) {
-		retryable, _ = DefaultRetryPolicy(context.Background(), resp, err)
-		retryAfter = DefaultBackoff(client.RetryWaitMin, client.RetryWaitMax, 1, resp)
-		return false, nil
-	}
+			client.CheckRetry = func(_ context.Context, resp *http.Response, err error) (bool, error) {
+				retryable, _ = DefaultRetryPolicy(context.Background(), resp, err)
+				retryAfter = DefaultBackoff(client.RetryWaitMin, client.RetryWaitMax, 1, resp)
+				return false, nil
+			}
 
-	_, err := client.Get(ts.URL)
-	if err != nil {
-		t.Fatalf("expected no errors since retryable")
-	}
+			_, err := client.Get(ts.URL)
+			if err != nil {
+				t.Fatalf("expected no errors since retryable")
+			}
 
-	if !retryable {
-		t.Fatal("Since 429 is recoverable, the default policy shall return true")
-	}
+			if !retryable {
+				t.Fatal("Since the error is recoverable, the default policy shall return true")
+			}
 
-	if retryAfter != 2*time.Second {
-		t.Fatalf("The header Retry-After specified 2 seconds, and shall not be %d seconds", retryAfter/time.Second)
+			if retryAfter != 2*time.Second {
+				t.Fatalf("The header Retry-After specified 2 seconds, and shall not be %d seconds", retryAfter/time.Second)
+			}
+		})
 	}
 }