Added jitter and added unit tests

Author: jprakash-db

src/databricks/sql/auth/retry.py

@@ -1,4 +1,5 @@
 import logging
+import random
 import time
 import typing
 from enum import Enum
@@ -287,14 +288,14 @@ def sleep_for_retry(self, response: BaseHTTPResponse) -> bool:
         if retry_after:
             proposed_wait = retry_after
         else:
-            proposed_wait = self.get_exponential_backoff()
+            proposed_wait = self.get_backoff_time()
 
         proposed_wait = min(proposed_wait, self.delay_max)
         self.check_proposed_wait(proposed_wait)
         time.sleep(proposed_wait)
         return True
 
-    def get_exponential_backoff(self) -> float:
+    def get_backoff_time(self) -> float:
         """
         This method implements the exponential backoff algorithm to calculate the delay between retries.
 
@@ -306,6 +307,9 @@ def get_exponential_backoff(self) -> float:
 
         current_attempt = self.stop_after_attempts_count - self.total
         proposed_backoff = (2**current_attempt) * self.delay_min
+        if self.backoff_jitter != 0.0:
+            proposed_backoff += random.random() * self.backoff_jitter
+
         proposed_backoff = min(proposed_backoff, self.delay_max)
         self.check_proposed_wait(proposed_backoff)
 

tests/unit/test_retry.py

@@ -1,11 +1,12 @@
 from os import error
 import time
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, patch, call
 import pytest
 from requests import Request
 from urllib3 import HTTPResponse
-from databricks.sql.auth.retry import DatabricksRetryPolicy, RequestHistory
-
+from databricks.sql.auth.retry import DatabricksRetryPolicy, RequestHistory, CommandType
+from databricks.sql.exc import MaxRetryDurationError
+from urllib3.exceptions import MaxRetryError
 
 class TestRetry:
     @pytest.fixture()
@@ -25,32 +26,55 @@ def error_history(self) -> RequestHistory:
             method="POST", url=None, error=None, status=503, redirect_location=None
         )
 
+    def calculate_backoff_time(self, attempt, delay_min, delay_max):
+        exponential_backoff_time = (2**attempt) * delay_min
+        return min(exponential_backoff_time, delay_max)
+
     @patch("time.sleep")
     def test_sleep__no_retry_after(self, t_mock, retry_policy, error_history):
         retry_policy._retry_start_time = time.time()
         retry_policy.history = [error_history, error_history]
         retry_policy.sleep(HTTPResponse(status=503))
-        t_mock.assert_called_with(2)
+
+        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")
-    def test_sleep__retry_after_is_binding(self, t_mock, retry_policy, error_history):
+    def test_sleep__no_retry_after_header__multiple_retries(self, t_mock, retry_policy):
+        num_attempts = retry_policy.stop_after_attempts_count
+
         retry_policy._retry_start_time = time.time()
-        retry_policy.history = [error_history, error_history]
-        retry_policy.sleep(HTTPResponse(status=503, headers={"Retry-After": "3"}))
-        t_mock.assert_called_with(3)
+        retry_policy.command_type = CommandType.OTHER
+
+        for attempt in range(num_attempts):
+            retry_policy.sleep(HTTPResponse(status=503))
+            # Internally urllib3 calls the increment function generating a new instance for every retry
+            retry_policy = retry_policy.increment()
+
+        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))
+
+        # 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])
 
     @patch("time.sleep")
-    def test_sleep__retry_after_present_but_not_binding(
-        self, t_mock, retry_policy, error_history
-    ):
+    def test_excessive_retry_attempts_error(self, t_mock, retry_policy):
+        # Attempting more than stop_after_attempt_count
+        num_attempts = retry_policy.stop_after_attempts_count + 1
+
         retry_policy._retry_start_time = time.time()
-        retry_policy.history = [error_history, error_history]
-        retry_policy.sleep(HTTPResponse(status=503, headers={"Retry-After": "1"}))
-        t_mock.assert_called_with(2)
+        retry_policy.command_type = CommandType.OTHER
+
+        with pytest.raises(MaxRetryError):
+            for attempt in range(num_attempts):
+                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_surpassed(self, t_mock, retry_policy, error_history):
+    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(4)
+        t_mock.assert_called_with(3)