SG-33891 Retries also on 504 (#316)

* Retries also on 504

* Use assertEqual instead

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

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):
         """