From 28b5efec0a1494cf2389530c40ee4c40b94eefdc Mon Sep 17 00:00:00 2001
From: Carlos Villavicencio
 <123113322+carlos-villavicencio-adsk@users.noreply.github.com>
Date: Tue, 19 Dec 2023 13:47:42 -0500
Subject: [PATCH] SG-33891 Retries also on 504 (#316)

* Retries also on 504

* Use assertEqual instead
---
 shotgun_api3/shotgun.py |  4 ++--
 tests/test_client.py    | 21 +++++++++++++++++++--
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/shotgun_api3/shotgun.py b/shotgun_api3/shotgun.py
index 50439f2a..9626056d 100644
--- a/shotgun_api3/shotgun.py
+++ b/shotgun_api3/shotgun.py
@@ -3443,8 +3443,8 @@ def _call_rpc(self, method, params, include_auth_params=True, first=False):
                 # We've seen some rare instances of SG returning 502 for issues that
                 # appear to be caused by something internal to SG. We're going to
                 # allow for limited retries for those specifically.
-                if attempt != max_attempts and e.errcode == 502:
-                    LOG.debug("Got a 502 response. Waiting and retrying...")
+                if attempt != max_attempts and e.errcode in [502, 504]:
+                    LOG.debug("Got a 502 or 504 response. Waiting and retrying...")
                     time.sleep(float(attempt) * backoff)
                     attempt += 1
                     continue
diff --git a/tests/test_client.py b/tests/test_client.py
index b1ccdbc3..95ceb4df 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -430,12 +430,29 @@ def test_call_rpc(self):
         expected = "rpc response with list result"
         self.assertEqual(d["results"], rv, expected)
 
-        # Test that we raise on a 502. This is ensuring the retries behavior
-        # in place specific to 502 responses still eventually ends up raising.
+        # Test that we raise on a 5xx. This is ensuring the retries behavior
+        # in place specific to 5xx responses still eventually ends up raising.
+        # 502
         d = {"results": ["foo", "bar"]}
         a = {"some": "args"}
         self._mock_http(d, status=(502, "bad gateway"))
         self.assertRaises(api.ProtocolError, self.sg._call_rpc, "list", a)
+        self.assertEqual(
+            4,
+            self.sg._http_request.call_count,
+            "Call is repeated up to 3 times",
+        )
+
+        # 504
+        d = {"results": ["foo", "bar"]}
+        a = {"some": "args"}
+        self._mock_http(d, status=(504, "gateway timeout"))
+        self.assertRaises(api.ProtocolError, self.sg._call_rpc, "list", a)
+        self.assertEqual(
+            4,
+            self.sg._http_request.call_count,
+            "Call is repeated up to 3 times",
+        )
 
     def test_upload_s3(self):
         """