From c1172aec592c8db24ce5220cbf5a077777458bc0 Mon Sep 17 00:00:00 2001
From: Cathy Ouyang <cathyo@google.com>
Date: Wed, 16 Mar 2022 13:55:40 -0700
Subject: [PATCH] fix: retry client side requests timeout

---
 .../requests/_request_helpers.py              |  1 +
 tests/unit/requests/test__helpers.py          | 34 ++++++++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/google/resumable_media/requests/_request_helpers.py b/google/resumable_media/requests/_request_helpers.py
index 5edfd7a5..fbeda5a3 100644
--- a/google/resumable_media/requests/_request_helpers.py
+++ b/google/resumable_media/requests/_request_helpers.py
@@ -37,6 +37,7 @@
 _CONNECTION_ERROR_CLASSES = (
     requests.exceptions.ConnectionError,
     requests.exceptions.ChunkedEncodingError,
+    requests.exceptions.Timeout,
     urllib3.exceptions.ProtocolError,
     ConnectionError,  # Python 3.x only, superclass of ConnectionResetError.
 )
diff --git a/tests/unit/requests/test__helpers.py b/tests/unit/requests/test__helpers.py
index 2d5d1950..5bd27cd8 100644
--- a/tests/unit/requests/test__helpers.py
+++ b/tests/unit/requests/test__helpers.py
@@ -200,7 +200,8 @@ def test_success_with_retry_connection_error(self, randint_mock, sleep_mock):
     def test_success_with_retry_chunked_encoding_error(self, randint_mock, sleep_mock):
         randint_mock.side_effect = [125, 625, 375]
 
-        response = _make_response(http.client.NOT_FOUND)
+        status_code = int(http.client.OK)
+        response = _make_response(status_code)
         responses = [
             requests.exceptions.ChunkedEncodingError,
             requests.exceptions.ChunkedEncodingError,
@@ -225,6 +226,37 @@ def test_success_with_retry_chunked_encoding_error(self, randint_mock, sleep_moc
         sleep_mock.assert_any_call(1.125)
         sleep_mock.assert_any_call(2.625)
 
+    @mock.patch(u"time.sleep")
+    @mock.patch(u"random.randint")
+    def test_success_with_retry_client_timeout(self, randint_mock, sleep_mock):
+        randint_mock.side_effect = [125, 625, 375]
+
+        status_code = int(http.client.OK)
+        response = _make_response(status_code)
+        responses = [
+            requests.exceptions.Timeout,
+            requests.exceptions.Timeout,
+            response,
+        ]
+        func = mock.Mock(side_effect=responses, spec=[])
+
+        retry_strategy = common.RetryStrategy()
+        ret_val = _request_helpers.wait_and_retry(
+            func, _get_status_code, retry_strategy
+        )
+
+        assert ret_val == responses[-1]
+
+        assert func.call_count == 3
+        assert func.mock_calls == [mock.call()] * 3
+
+        assert randint_mock.call_count == 2
+        assert randint_mock.mock_calls == [mock.call(0, 1000)] * 2
+
+        assert sleep_mock.call_count == 2
+        sleep_mock.assert_any_call(1.125)
+        sleep_mock.assert_any_call(2.625)
+
     @mock.patch("time.sleep")
     @mock.patch("random.randint")
     def test_retry_exceeds_max_cumulative(self, randint_mock, sleep_mock):