Skip to content

Commit

Permalink
Fixed the retry tests failing (#500)
Browse files Browse the repository at this point in the history
  • Loading branch information
jprakash-db authored Feb 2, 2025
1 parent fc9da22 commit accf6ff
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 21 deletions.
20 changes: 10 additions & 10 deletions tests/e2e/common/retry_test_mixins.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,9 +121,9 @@ class PySQLRetryTestsMixin:
# For testing purposes
_retry_policy = {
"_retry_delay_min": 0.1,
"_retry_delay_max": 5,
"_retry_delay_max": 3,
"_retry_stop_after_attempts_count": 5,
"_retry_stop_after_attempts_duration": 10,
"_retry_stop_after_attempts_duration": 30,
"_retry_delay_default": 0.5,
}

Expand All @@ -135,7 +135,7 @@ def test_retry_urllib3_settings_are_honored(self):
urllib3_config = {"connect": 10, "read": 11, "redirect": 12}
rp = DatabricksRetryPolicy(
delay_min=0.1,
delay_max=10.0,
delay_max=3,
stop_after_attempts_count=10,
stop_after_attempts_duration=10.0,
delay_default=1.0,
Expand Down Expand Up @@ -174,14 +174,14 @@ def test_retry_max_count_not_exceeded(self):
def test_retry_exponential_backoff(self):
"""GIVEN the retry policy is configured for reasonable exponential backoff
WHEN the server sends nothing but 429 responses with retry-afters
THEN the connector will use those retry-afters values as delay
THEN the connector will use those retry-afters values as floor
"""
retry_policy = self._retry_policy.copy()
retry_policy["_retry_delay_min"] = 1

time_start = time.time()
with mocked_server_response(
status=429, headers={"Retry-After": "3"}
status=429, headers={"Retry-After": "8"}
) as mock_obj:
with pytest.raises(RequestError) as cm:
with self.connection(extra_params=retry_policy) as conn:
Expand All @@ -191,14 +191,14 @@ def test_retry_exponential_backoff(self):
assert isinstance(cm.value.args[1], MaxRetryDurationError)

# With setting delay_min to 1, the expected retry delays should be:
# 3, 3, 3, 3
# 8, 8, 8, 8
# The first 3 retries are allowed, the 4th retry puts the total duration over the limit
# of 10 seconds
# of 30 seconds
assert mock_obj.return_value.getresponse.call_count == 4
assert duration > 6
assert duration > 24

# Should be less than 7, but this is a safe margin for CI/CD slowness
assert duration < 10
# Should be less than 26, but this is a safe margin for CI/CD slowness
assert duration < 30

def test_retry_max_duration_not_exceeded(self):
"""GIVEN the max attempt duration of 10 seconds
Expand Down
21 changes: 10 additions & 11 deletions tests/unit/test_retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,11 @@ 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 = max(
self.calculate_backoff_time(
0, retry_policy.delay_min, retry_policy.delay_max
),
retry_policy.delay_max,
)
t_mock.assert_called_with(expected_backoff_time)

Expand All @@ -54,8 +57,11 @@ 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
max(
self.calculate_backoff_time(
attempt, retry_policy.delay_min, retry_policy.delay_max
),
retry_policy.delay_max,
)
)

Expand All @@ -77,10 +83,3 @@ def test_excessive_retry_attempts_error(self, t_mock, retry_policy):
retry_policy.sleep(HTTPResponse(status=503))
# Internally urllib3 calls the increment function generating a new instance for every retry
retry_policy = retry_policy.increment()

@patch("time.sleep")
def test_sleep__retry_after_present(self, t_mock, retry_policy, error_history):
retry_policy._retry_start_time = time.time()
retry_policy.history = [error_history, error_history, error_history]
retry_policy.sleep(HTTPResponse(status=503, headers={"Retry-After": "3"}))
t_mock.assert_called_with(3)

0 comments on commit accf6ff

Please sign in to comment.