SG-34197 Retry S3 uploads on error 500 (#324)

* Retry S3 uploads on error 500

* Remove context manager variable

* Improve conditional logic

Author: carlos-villavicencio-adsk

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:

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.",
                                                      {},

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