5
5
from databricks .sql .auth .retry import DatabricksRetryPolicy , RequestHistory , CommandType
6
6
from urllib3 .exceptions import MaxRetryError
7
7
8
+
8
9
class TestRetry :
9
10
@pytest .fixture ()
10
11
def retry_policy (self ) -> DatabricksRetryPolicy :
@@ -33,7 +34,9 @@ def test_sleep__no_retry_after(self, t_mock, retry_policy, error_history):
33
34
retry_policy .history = [error_history , error_history ]
34
35
retry_policy .sleep (HTTPResponse (status = 503 ))
35
36
36
- expected_backoff_time = self .calculate_backoff_time (0 , retry_policy .delay_min , retry_policy .delay_max )
37
+ expected_backoff_time = self .calculate_backoff_time (
38
+ 0 , retry_policy .delay_min , retry_policy .delay_max
39
+ )
37
40
t_mock .assert_called_with (expected_backoff_time )
38
41
39
42
@patch ("time.sleep" )
@@ -50,10 +53,16 @@ def test_sleep__no_retry_after_header__multiple_retries(self, t_mock, retry_poli
50
53
51
54
expected_backoff_times = []
52
55
for attempt in range (num_attempts ):
53
- expected_backoff_times .append (self .calculate_backoff_time (attempt , retry_policy .delay_min , retry_policy .delay_max ))
56
+ expected_backoff_times .append (
57
+ self .calculate_backoff_time (
58
+ attempt , retry_policy .delay_min , retry_policy .delay_max
59
+ )
60
+ )
54
61
55
62
# Asserts if the sleep value was called in the expected order
56
- t_mock .assert_has_calls ([call (expected_time ) for expected_time in expected_backoff_times ])
63
+ t_mock .assert_has_calls (
64
+ [call (expected_time ) for expected_time in expected_backoff_times ]
65
+ )
57
66
58
67
@patch ("time.sleep" )
59
68
def test_excessive_retry_attempts_error (self , t_mock , retry_policy ):
0 commit comments