Fix for check_types github action failing (#472)

Fixed the chekc_types failing

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: jprakash-db

src/databricks/sql/auth/retry.py

@@ -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

tests/unit/test_retry.py

@@ -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:
@@ -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")
@@ -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):