refactor(client): move retry logic into utility

nejch · max-wittig · commit 3235c48328c2 · 2024-09-10T17:07:18.000+02:00
diff --git a/gitlab/client.py b/gitlab/client.py
@@ -2,7 +2,6 @@
 
 import os
 import re
-import time
 from typing import (
     Any,
     BinaryIO,
@@ -718,7 +717,12 @@ def http_request(
         send_data = self._backend.prepare_send_data(files, post_data, raw)
         opts["headers"]["Content-type"] = send_data.content_type
 
-        cur_retries = 0
+        retry = utils.Retry(
+            max_retries=max_retries,
+            obey_rate_limit=obey_rate_limit,
+            retry_transient_errors=retry_transient_errors,
+        )
+
         while True:
             try:
                 result = self._backend.http_request(
@@ -733,46 +737,17 @@ def http_request(
                     **opts,
                 )
             except (requests.ConnectionError, requests.exceptions.ChunkedEncodingError):
-                if retry_transient_errors and (
-                    max_retries == -1 or cur_retries < max_retries
-                ):
-                    wait_time = 2**cur_retries * 0.1
-                    cur_retries += 1
-                    time.sleep(wait_time)
+                if retry.handle_retry():
                     continue
-
                 raise
 
             self._check_redirects(result.response)
 
             if 200 <= result.status_code < 300:
                 return result.response
 
-            def should_retry() -> bool:
-                if result.status_code == 429 and obey_rate_limit:
-                    return True
-
-                if not retry_transient_errors:
-                    return False
-                if result.status_code in gitlab.const.RETRYABLE_TRANSIENT_ERROR_CODES:
-                    return True
-                if result.status_code == 409 and "Resource lock" in result.reason:
-                    return True
-
-                return False
-
-            if should_retry():
-                # Response headers documentation:
-                # https://docs.gitlab.com/ee/user/admin_area/settings/user_and_ip_rate_limits.html#response-headers
-                if max_retries == -1 or cur_retries < max_retries:
-                    wait_time = 2**cur_retries * 0.1
-                    if "Retry-After" in result.headers:
-                        wait_time = int(result.headers["Retry-After"])
-                    elif "RateLimit-Reset" in result.headers:
-                        wait_time = int(result.headers["RateLimit-Reset"]) - time.time()
-                    cur_retries += 1
-                    time.sleep(wait_time)
-                    continue
+            if retry.handle_retry_on_status(result):
+                continue
 
             error_message = result.content
             try:
diff --git a/gitlab/utils.py b/gitlab/utils.py
@@ -2,6 +2,7 @@
 import email.message
 import logging
 import pathlib
+import time
 import traceback
 import urllib.parse
 import warnings
@@ -10,6 +11,7 @@
 import requests
 
 from gitlab import const, types
+from gitlab._backends import requests_backend
 
 
 class _StdoutStream:
@@ -85,6 +87,64 @@ def response_content(
     return None
 
 
+class Retry:
+    def __init__(
+        self,
+        max_retries: int,
+        obey_rate_limit: Optional[bool] = True,
+        retry_transient_errors: Optional[bool] = False,
+    ) -> None:
+        self.cur_retries = 0
+        self.max_retries = max_retries
+        self.obey_rate_limit = obey_rate_limit
+        self.retry_transient_errors = retry_transient_errors
+
+    def _retryable_status_code(
+        self,
+        result: requests_backend.RequestsResponse,
+    ) -> bool:
+        if result.status_code == 429 and self.obey_rate_limit:
+            return True
+
+        if not self.retry_transient_errors:
+            return False
+        if result.status_code in const.RETRYABLE_TRANSIENT_ERROR_CODES:
+            return True
+        if result.status_code == 409 and "Resource lock" in result.reason:
+            return True
+
+        return False
+
+    def handle_retry_on_status(self, result: requests_backend.RequestsResponse) -> bool:
+        if not self._retryable_status_code(result):
+            return False
+
+        # Response headers documentation:
+        # https://docs.gitlab.com/ee/user/admin_area/settings/user_and_ip_rate_limits.html#response-headers
+        if self.max_retries == -1 or self.cur_retries < self.max_retries:
+            wait_time = 2**self.cur_retries * 0.1
+            if "Retry-After" in result.headers:
+                wait_time = int(result.headers["Retry-After"])
+            elif "RateLimit-Reset" in result.headers:
+                wait_time = int(result.headers["RateLimit-Reset"]) - time.time()
+            self.cur_retries += 1
+            time.sleep(wait_time)
+            return True
+
+        return False
+
+    def handle_retry(self) -> bool:
+        if self.retry_transient_errors and (
+            self.max_retries == -1 or self.cur_retries < self.max_retries
+        ):
+            wait_time = 2**self.cur_retries * 0.1
+            self.cur_retries += 1
+            time.sleep(wait_time)
+            return True
+
+        return False
+
+
 def _transform_types(
     data: Dict[str, Any],
     custom_types: Dict[str, Any],
diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py
@@ -0,0 +1,58 @@
+import time
+from unittest import mock
+
+import pytest
+import requests
+
+from gitlab import utils
+from gitlab._backends import requests_backend
+
+
+def test_handle_retry_on_status_ignores_unknown_status_code():
+    retry = utils.Retry(max_retries=1, retry_transient_errors=True)
+    response = requests.Response()
+    response.status_code = 418
+    backend_response = requests_backend.RequestsResponse(response)
+
+    assert retry.handle_retry_on_status(backend_response) is False
+
+
+def test_handle_retry_on_status_accepts_retry_after_header(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    mock_sleep = mock.Mock()
+    monkeypatch.setattr(time, "sleep", mock_sleep)
+
+    retry = utils.Retry(max_retries=1)
+    response = requests.Response()
+    response.status_code = 429
+    response.headers["Retry-After"] = "1"
+    backend_response = requests_backend.RequestsResponse(response)
+
+    assert retry.handle_retry_on_status(backend_response) is True
+    assert isinstance(mock_sleep.call_args[0][0], int)
+
+
+def test_handle_retry_on_status_accepts_ratelimit_reset_header(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    mock_sleep = mock.Mock()
+    monkeypatch.setattr(time, "sleep", mock_sleep)
+
+    retry = utils.Retry(max_retries=1)
+    response = requests.Response()
+    response.status_code = 429
+    response.headers["RateLimit-Reset"] = str(int(time.time() + 1))
+    backend_response = requests_backend.RequestsResponse(response)
+
+    assert retry.handle_retry_on_status(backend_response) is True
+    assert isinstance(mock_sleep.call_args[0][0], float)
+
+
+def test_handle_retry_on_status_returns_false_when_max_retries_reached():
+    retry = utils.Retry(max_retries=0)
+    response = requests.Response()
+    response.status_code = 429
+    backend_response = requests_backend.RequestsResponse(response)
+
+    assert retry.handle_retry_on_status(backend_response) is False
