From a1c5f7d0e3ce48d8d6eb8aced31707a881f7ee96 Mon Sep 17 00:00:00 2001
From: cojenco <cathyo@google.com>
Date: Mon, 30 Aug 2021 14:32:55 -0700
Subject: [PATCH] fix: do not mark upload download instances invalid with
 retriable error codes (#261)

---
 google/resumable_media/_helpers.py |  3 ++-
 tests/unit/test__helpers.py        | 22 ++++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/google/resumable_media/_helpers.py b/google/resumable_media/_helpers.py
index 5f45ad5c..a105d3c6 100644
--- a/google/resumable_media/_helpers.py
+++ b/google/resumable_media/_helpers.py
@@ -95,7 +95,8 @@ def require_status_code(response, status_codes, get_status_code, callback=do_not
     """
     status_code = get_status_code(response)
     if status_code not in status_codes:
-        callback()
+        if status_code not in common.RETRYABLE:
+            callback()
         raise common.InvalidResponse(
             response,
             "Request failed with status code",
diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py
index 62a527ca..498ca290 100644
--- a/tests/unit/test__helpers.py
+++ b/tests/unit/test__helpers.py
@@ -125,6 +125,28 @@ def test_failure_with_callback(self):
         assert error.args[3:] == status_codes
         callback.assert_called_once_with()
 
+    def test_retryable_failure_without_callback(self):
+        status_codes = (http.client.OK,)
+        retryable_responses = [
+            _make_response(status_code) for status_code in common.RETRYABLE
+        ]
+        callback = mock.Mock(spec=[])
+        for retryable_response in retryable_responses:
+            with pytest.raises(common.InvalidResponse) as exc_info:
+                _helpers.require_status_code(
+                    retryable_response,
+                    status_codes,
+                    self._get_status_code,
+                    callback=callback,
+                )
+
+            error = exc_info.value
+            assert error.response is retryable_response
+            assert len(error.args) == 4
+            assert error.args[1] == retryable_response.status_code
+            assert error.args[3:] == status_codes
+            callback.assert_not_called()
+
 
 class Test_calculate_retry_wait(object):
     @mock.patch("random.randint", return_value=125)