diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py index ec321bed..0243d0aa 100755 --- a/src/databricks/sql/auth/retry.py +++ b/src/databricks/sql/auth/retry.py @@ -305,7 +305,7 @@ def get_backoff_time(self) -> float: :return: """ - current_attempt = self.stop_after_attempts_count - self.total + current_attempt = self.stop_after_attempts_count - int(self.total or 0) proposed_backoff = (2**current_attempt) * self.delay_min if self.backoff_jitter != 0.0: proposed_backoff += random.random() * self.backoff_jitter diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py index b7648ffb..1e18e1f4 100644 --- a/tests/unit/test_retry.py +++ b/tests/unit/test_retry.py @@ -5,6 +5,7 @@ from databricks.sql.auth.retry import DatabricksRetryPolicy, RequestHistory, CommandType from urllib3.exceptions import MaxRetryError + class TestRetry: @pytest.fixture() def retry_policy(self) -> DatabricksRetryPolicy: @@ -33,7 +34,9 @@ def test_sleep__no_retry_after(self, t_mock, retry_policy, error_history): retry_policy.history = [error_history, error_history] retry_policy.sleep(HTTPResponse(status=503)) - expected_backoff_time = self.calculate_backoff_time(0, retry_policy.delay_min, retry_policy.delay_max) + expected_backoff_time = self.calculate_backoff_time( + 0, retry_policy.delay_min, retry_policy.delay_max + ) t_mock.assert_called_with(expected_backoff_time) @patch("time.sleep") @@ -50,10 +53,16 @@ def test_sleep__no_retry_after_header__multiple_retries(self, t_mock, retry_poli expected_backoff_times = [] for attempt in range(num_attempts): - expected_backoff_times.append(self.calculate_backoff_time(attempt, retry_policy.delay_min, retry_policy.delay_max)) + expected_backoff_times.append( + self.calculate_backoff_time( + attempt, retry_policy.delay_min, retry_policy.delay_max + ) + ) # Asserts if the sleep value was called in the expected order - t_mock.assert_has_calls([call(expected_time) for expected_time in expected_backoff_times]) + t_mock.assert_has_calls( + [call(expected_time) for expected_time in expected_backoff_times] + ) @patch("time.sleep") def test_excessive_retry_attempts_error(self, t_mock, retry_policy):