fix unit

Signed-off-by: Vikrant Puppala <[email protected]>

Author: vikrantpuppala

src/databricks/sql/auth/common.py

@@ -72,7 +72,9 @@ def __init__(
         self.retry_stop_after_attempts_count = retry_stop_after_attempts_count or 5
         self.retry_delay_min = retry_delay_min or 1.0
         self.retry_delay_max = retry_delay_max or 60.0
-        self.retry_stop_after_attempts_duration = retry_stop_after_attempts_duration or 900.0
+        self.retry_stop_after_attempts_duration = (
+            retry_stop_after_attempts_duration or 900.0
+        )
         self.retry_delay_default = retry_delay_default or 5.0
         self.retry_dangerous_codes = retry_dangerous_codes or []
         self.http_proxy = http_proxy

src/databricks/sql/client.py

@@ -376,10 +376,14 @@ def _build_client_context(self, server_hostname: str, **kwargs):
             hostname=server_hostname,
             ssl_options=ssl_options,
             socket_timeout=kwargs.get("_socket_timeout"),
-            retry_stop_after_attempts_count=kwargs.get("_retry_stop_after_attempts_count"),
+            retry_stop_after_attempts_count=kwargs.get(
+                "_retry_stop_after_attempts_count"
+            ),
             retry_delay_min=kwargs.get("_retry_delay_min"),
             retry_delay_max=kwargs.get("_retry_delay_max"),
-            retry_stop_after_attempts_duration=kwargs.get("_retry_stop_after_attempts_duration"),
+            retry_stop_after_attempts_duration=kwargs.get(
+                "_retry_stop_after_attempts_duration"
+            ),
             retry_delay_default=kwargs.get("_retry_delay_default"),
             retry_dangerous_codes=kwargs.get("_retry_dangerous_codes"),
             http_proxy=kwargs.get("_http_proxy"),

src/databricks/sql/common/unified_http_client.py

@@ -78,7 +78,7 @@ def _setup_pool_manager(self):
             delay_default=self.config.retry_delay_default,
             force_dangerous_codes=self.config.retry_dangerous_codes,
         )
-        
+
         # Initialize the required attributes that DatabricksRetryPolicy expects
         # but doesn't initialize in its constructor
         self._retry_policy._command_type = None
@@ -152,10 +152,10 @@ def request_context(
         logger.debug("Making %s request to %s", method, url)
 
         request_headers = self._prepare_headers(headers)
-        
+
         # Prepare retry policy for this request
         self._prepare_retry_policy()
-        
+
         response = None
 
         try:

tests/unit/test_retry.py

@@ -34,7 +34,7 @@ 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 = max(
+        expected_backoff_time = min(
             self.calculate_backoff_time(
                 0, retry_policy.delay_min, retry_policy.delay_max
             ),
@@ -57,7 +57,7 @@ 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(
-                max(
+                min(
                     self.calculate_backoff_time(
                         attempt, retry_policy.delay_min, retry_policy.delay_max
                     ),