fix: check if retry is allowed after retry wait calculation (#258)

* fix: check if retry is allowed after retry wait calculation

* revise docstrings and add test case

google/resumable_media/_helpers.py

@@ -192,17 +192,17 @@ def wait_and_retry(func, get_status_code, retry_strategy):
         else:
             return response
 
-        if not retry_strategy.retry_allowed(total_sleep, num_retries):
-            # Retries are exhausted and no acceptable response was received.
-            # Raise the retriable_error.
-            raise error
-
         base_wait, wait_time = calculate_retry_wait(
             base_wait, retry_strategy.max_sleep, retry_strategy.multiplier
         )
-
         num_retries += 1
         total_sleep += wait_time
+
+        # Check if (another) retry is allowed. If retries are exhausted and
+        # no acceptable response was received, raise the retriable error.
+        if not retry_strategy.retry_allowed(total_sleep, num_retries):
+            raise error
+
         time.sleep(wait_time)
 
 

google/resumable_media/common.py

@@ -158,9 +158,10 @@ def retry_allowed(self, total_sleep, num_retries):
         """Check if another retry is allowed.
 
         Args:
-            total_sleep (float): The amount of sleep accumulated by the caller.
-            num_retries (int): The number of retries already attempted by
-                the caller.
+            total_sleep (float): With another retry, the amount of sleep that
+                will be accumulated by the caller.
+            num_retries (int): With another retry, the number of retries that
+                will be attempted by the caller.
 
         Returns:
             bool: Indicating if another retry is allowed (depending on either

tests/unit/test__helpers.py

@@ -337,7 +337,7 @@ def test_connection_import_error_failure(self, randint_mock, sleep_mock):
     @mock.patch("time.sleep")
     @mock.patch("random.randint")
     def test_retry_exceeds_max_cumulative(self, randint_mock, sleep_mock):
-        randint_mock.side_effect = [875, 0, 375, 500, 500, 250, 125, 0]
+        randint_mock.side_effect = [875, 0, 375, 500, 500, 250, 125]
 
         status_codes = (
             http.client.SERVICE_UNAVAILABLE,
@@ -346,7 +346,6 @@ def test_retry_exceeds_max_cumulative(self, randint_mock, sleep_mock):
             http.client.INTERNAL_SERVER_ERROR,
             http.client.SERVICE_UNAVAILABLE,
             http.client.BAD_GATEWAY,
-            http.client.GATEWAY_TIMEOUT,
             common.TOO_MANY_REQUESTS,
         )
         responses = [_make_response(status_code) for status_code in status_codes]
@@ -365,47 +364,119 @@ def raise_response():
         assert ret_val.status_code == status_codes[-1]
         assert status_codes[-1] in common.RETRYABLE
 
-        assert func.call_count == 8
-        assert func.mock_calls == [mock.call()] * 8
+        assert func.call_count == 7
+        assert func.mock_calls == [mock.call()] * 7
 
         assert randint_mock.call_count == 7
         assert randint_mock.mock_calls == [mock.call(0, 1000)] * 7
 
-        assert sleep_mock.call_count == 7
+        assert sleep_mock.call_count == 6
         sleep_mock.assert_any_call(1.875)
         sleep_mock.assert_any_call(2.0)
         sleep_mock.assert_any_call(4.375)
         sleep_mock.assert_any_call(8.5)
         sleep_mock.assert_any_call(16.5)
         sleep_mock.assert_any_call(32.25)
-        sleep_mock.assert_any_call(64.125)
+
+    @mock.patch("time.sleep")
+    @mock.patch("random.randint")
+    def test_retry_exceeds_max_retries(self, randint_mock, sleep_mock):
+        randint_mock.side_effect = [875, 0, 375, 500, 500, 250, 125]
+
+        status_codes = (
+            http.client.SERVICE_UNAVAILABLE,
+            http.client.GATEWAY_TIMEOUT,
+            common.TOO_MANY_REQUESTS,
+            http.client.INTERNAL_SERVER_ERROR,
+            http.client.SERVICE_UNAVAILABLE,
+            http.client.BAD_GATEWAY,
+            common.TOO_MANY_REQUESTS,
+        )
+        responses = [_make_response(status_code) for status_code in status_codes]
+
+        def raise_response():
+            raise common.InvalidResponse(responses.pop(0))
+
+        func = mock.Mock(side_effect=raise_response)
+
+        retry_strategy = common.RetryStrategy(max_retries=6)
+        try:
+            _helpers.wait_and_retry(func, _get_status_code, retry_strategy)
+        except common.InvalidResponse as e:
+            ret_val = e.response
+
+        assert ret_val.status_code == status_codes[-1]
+        assert status_codes[-1] in common.RETRYABLE
+
+        assert func.call_count == 7
+        assert func.mock_calls == [mock.call()] * 7
+
+        assert randint_mock.call_count == 7
+        assert randint_mock.mock_calls == [mock.call(0, 1000)] * 7
+
+        assert sleep_mock.call_count == 6
+        sleep_mock.assert_any_call(1.875)
+        sleep_mock.assert_any_call(2.0)
+        sleep_mock.assert_any_call(4.375)
+        sleep_mock.assert_any_call(8.5)
+        sleep_mock.assert_any_call(16.5)
+        sleep_mock.assert_any_call(32.25)
+
+    @mock.patch("time.sleep")
+    @mock.patch("random.randint")
+    def test_retry_zero_max_retries(self, randint_mock, sleep_mock):
+        randint_mock.side_effect = [875, 0, 375]
+
+        status_codes = (
+            http.client.SERVICE_UNAVAILABLE,
+            http.client.GATEWAY_TIMEOUT,
+            common.TOO_MANY_REQUESTS,
+        )
+        responses = [_make_response(status_code) for status_code in status_codes]
+
+        def raise_response():
+            raise common.InvalidResponse(responses.pop(0))
+
+        func = mock.Mock(side_effect=raise_response)
+
+        retry_strategy = common.RetryStrategy(max_retries=0)
+        try:
+            _helpers.wait_and_retry(func, _get_status_code, retry_strategy)
+        except common.InvalidResponse as e:
+            ret_val = e.response
+
+        assert func.call_count == 1
+        assert func.mock_calls == [mock.call()] * 1
+        assert ret_val.status_code == status_codes[0]
+
+        assert randint_mock.call_count == 1
+        assert sleep_mock.call_count == 0
 
     @mock.patch("time.sleep")
     @mock.patch("random.randint")
     def test_retry_exceeded_reraises_connection_error(self, randint_mock, sleep_mock):
         randint_mock.side_effect = [875, 0, 375, 500, 500, 250, 125]
 
-        responses = [requests.exceptions.ConnectionError] * 8
+        responses = [requests.exceptions.ConnectionError] * 7
         func = mock.Mock(side_effect=responses, spec=[])
 
         retry_strategy = common.RetryStrategy(max_cumulative_retry=100.0)
         with pytest.raises(requests.exceptions.ConnectionError):
             _helpers.wait_and_retry(func, _get_status_code, retry_strategy)
 
-        assert func.call_count == 8
-        assert func.mock_calls == [mock.call()] * 8
+        assert func.call_count == 7
+        assert func.mock_calls == [mock.call()] * 7
 
         assert randint_mock.call_count == 7
         assert randint_mock.mock_calls == [mock.call(0, 1000)] * 7
 
-        assert sleep_mock.call_count == 7
+        assert sleep_mock.call_count == 6
         sleep_mock.assert_any_call(1.875)
         sleep_mock.assert_any_call(2.0)
         sleep_mock.assert_any_call(4.375)
         sleep_mock.assert_any_call(8.5)
         sleep_mock.assert_any_call(16.5)
         sleep_mock.assert_any_call(32.25)
-        sleep_mock.assert_any_call(64.125)
 
 
 def _make_response(status_code):