Retry rate limits on chunk uploading

lichray · lichray · commit cc6d398dfd9e · 2016-07-28T13:45:22.000-05:00
diff --git a/googleapiclient/http.py b/googleapiclient/http.py
@@ -639,7 +639,7 @@ def next_chunk(self, num_retries=0):
     """Get the next chunk of the download.
 
     Args:
-      num_retries: Integer, number of times to retry 500's with randomized
+      num_retries: Integer, number of times to retry with randomized
             exponential backoff. If all retries fail, the raised HttpError
             represents the last request. If zero (default), we attempt the
             request only once.
@@ -782,7 +782,7 @@ def execute(self, http=None, num_retries=0):
     Args:
       http: httplib2.Http, an http object to be used in place of the
             one the HttpRequest request object was constructed with.
-      num_retries: Integer, number of times to retry 500's with randomized
+      num_retries: Integer, number of times to retry with randomized
             exponential backoff. If all retries fail, the raised HttpError
             represents the last request. If zero (default), we attempt the
             request only once.
@@ -870,7 +870,7 @@ def next_chunk(self, http=None, num_retries=0):
     Args:
       http: httplib2.Http, an http object to be used in place of the
             one the HttpRequest request object was constructed with.
-      num_retries: Integer, number of times to retry 500's with randomized
+      num_retries: Integer, number of times to retry with randomized
             exponential backoff. If all retries fail, the raised HttpError
             represents the last request. If zero (default), we attempt the
             request only once.
@@ -965,7 +965,7 @@ def next_chunk(self, http=None, num_retries=0):
       except:
         self._in_error_state = True
         raise
-      if resp.status < 500:
+      if not _should_retry_response(resp.status, content):
         break
 
     return self._process_response(resp, content)
diff --git a/tests/test_http.py b/tests/test_http.py
@@ -341,16 +341,16 @@ def test_media_io_base_next_chunk_retries(self):
     upload = MediaIoBaseUpload(
         fd=fd, mimetype='image/png', chunksize=500, resumable=True)
 
-    # Simulate 5XXs for both the request that creates the resumable upload and
-    # the upload itself.
+    # Simulate errors for both the request that creates the resumable upload
+    # and the upload itself.
     http = HttpMockSequence([
       ({'status': '500'}, ''),
       ({'status': '500'}, ''),
       ({'status': '503'}, ''),
       ({'status': '200', 'location': 'location'}, ''),
-      ({'status': '500'}, ''),
-      ({'status': '500'}, ''),
-      ({'status': '503'}, ''),
+      ({'status': '403'}, USER_RATE_LIMIT_EXCEEDED_RESPONSE),
+      ({'status': '403'}, RATE_LIMIT_EXCEEDED_RESPONSE),
+      ({'status': '429'}, ''),
       ({'status': '200'}, '{}'),
     ])
 
@@ -372,6 +372,34 @@ def test_media_io_base_next_chunk_retries(self):
     request.execute(num_retries=3)
     self.assertEqual([20, 40, 80, 20, 40, 80], sleeptimes)
 
+  def test_media_io_base_next_chunk_no_retry_403_not_configured(self):
+    fd = BytesIO(b"i am png")
+    upload = MediaIoBaseUpload(
+        fd=fd, mimetype='image/png', chunksize=500, resumable=True)
+
+    http = HttpMockSequence([
+        ({'status': '403'}, NOT_CONFIGURED_RESPONSE),
+        ({'status': '200'}, '{}')
+        ])
+
+    model = JsonModel()
+    uri = u'https://www.googleapis.com/someapi/v1/upload/?foo=bar'
+    method = u'POST'
+    request = HttpRequest(
+        http,
+        model.response,
+        uri,
+        method=method,
+        headers={},
+        resumable=upload)
+
+    request._rand = lambda: 1.0
+    request._sleep =  mock.MagicMock()
+
+    with self.assertRaises(HttpError):
+      request.execute(num_retries=3)
+    request._sleep.assert_not_called()
+
 
 class TestMediaIoBaseDownload(unittest.TestCase):
 
