Skip to content

Commit

Permalink
Fix for check_types github action failing (#472)
Browse files Browse the repository at this point in the history
Fixed the chekc_types failing
  • Loading branch information
jprakash-db authored Nov 26, 2024
1 parent 328aeb5 commit 980af88
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 4 deletions.
2 changes: 1 addition & 1 deletion src/databricks/sql/auth/retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 12 additions & 3 deletions tests/unit/test_retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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")
Expand All @@ -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):
Expand Down

0 comments on commit 980af88

Please sign in to comment.