From d504a5aa63ec3ca688466e14951257a0eb89cdd4 Mon Sep 17 00:00:00 2001
From: Carlos Villavicencio
 <123113322+carlos-villavicencio-adsk@users.noreply.github.com>
Date: Mon, 29 Jan 2024 08:11:05 -0500
Subject: [PATCH] SG-34197 Retry S3 uploads on error 500 (#324)

* Retry S3 uploads on error 500

* Remove context manager variable

* Improve conditional logic
---
 shotgun_api3/shotgun.py | 10 ++++------
 tests/base.py           |  4 ++--
 tests/test_client.py    | 26 +++++++++++++++++++++++++-
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/shotgun_api3/shotgun.py b/shotgun_api3/shotgun.py
index 20071708..a3b78974 100644
--- a/shotgun_api3/shotgun.py
+++ b/shotgun_api3/shotgun.py
@@ -4145,16 +4145,14 @@ def _upload_data_to_storage(self, data, content_type, size, storage_url):
                 LOG.debug("Completed request to %s" % request.get_method())
 
             except urllib.error.HTTPError as e:
-                if e.code == 500:
-                    raise ShotgunError("Server encountered an internal error.\n%s\n%s\n\n" % (storage_url, e))
-                elif attempt != max_attempts and e.code == 503:
-                    LOG.debug("Got a 503 response. Waiting and retrying...")
+                if attempt != max_attempts and e.code in [500, 503]:
+                    LOG.debug("Got a %s response. Waiting and retrying..." % e.code)
                     time.sleep(float(attempt) * backoff)
                     attempt += 1
                     continue
+                elif e.code in [500, 503]:
+                    raise ShotgunError("Got a %s response when uploading to %s: %s" % (e.code, storage_url, e))
                 else:
-                    if e.code == 503:
-                        raise ShotgunError("Got a 503 response when uploading to %s: %s" % (storage_url, e))
                     raise ShotgunError("Unanticipated error occurred uploading to %s: %s" % (storage_url, e))
 
             else:
diff --git a/tests/base.py b/tests/base.py
index 8296061a..5f4b61aa 100644
--- a/tests/base.py
+++ b/tests/base.py
@@ -119,7 +119,7 @@ def setUp(self):
         self._setup_mock()
         self._setup_mock_data()
 
-    def _setup_mock(self):
+    def _setup_mock(self, s3_status_code_error=503):
         """Setup mocking on the ShotgunClient to stop it calling a live server
         """
         # Replace the function used to make the final call to the server
@@ -131,7 +131,7 @@ def _setup_mock(self):
         self.sg._make_upload_request = mock.Mock(spec=api.Shotgun._make_upload_request,
                                                  side_effect = urllib.error.HTTPError(
                                                      "url",
-                                                     503,
+                                                     s3_status_code_error,
                                                      "The server is currently down or to busy to reply."
                                                      "Please try again later.",
                                                      {},
diff --git a/tests/test_client.py b/tests/test_client.py
index 95ceb4df..f7e783b2 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -454,7 +454,7 @@ def test_call_rpc(self):
             "Call is repeated up to 3 times",
         )
 
-    def test_upload_s3(self):
+    def test_upload_s3_503(self):
         """
         Test 503 response is retried when uploading to S3.
         """
@@ -476,6 +476,30 @@ def test_upload_s3(self):
         self.assertTrue(
             max_attempts == self.sg._make_upload_request.call_count,
             "Call is repeated up to 3 times")
+    
+    def test_upload_s3_500(self):
+        """
+        Test 500 response is retried when uploading to S3.
+        """
+        self._setup_mock(s3_status_code_error=500)
+        this_dir, _ = os.path.split(__file__)
+        storage_url = "http://foo.com/"
+        path = os.path.abspath(os.path.expanduser(
+            os.path.join(this_dir, "sg_logo.jpg")))
+        max_attempts = 4  # Max retries to S3 server attempts
+        # Expected HTTPError exception error message
+        expected = "The server is currently down or to busy to reply." \
+                   "Please try again later."
+
+        # Test the Internal function that is used to upload each
+        # data part in the context of multi-part uploads to S3, we
+        # simulate the HTTPError exception raised with 503 status errors
+        with self.assertRaises(api.ShotgunError, msg=expected):
+            self.sg._upload_file_to_storage(path, storage_url)
+        # Test the max retries attempt
+        self.assertTrue(
+            max_attempts == self.sg._make_upload_request.call_count,
+            "Call is repeated up to 3 times")
 
     def test_transform_data(self):
         """Outbound data is transformed"""