Automatic retry logic

onyx/core.py

@@ -3,6 +3,7 @@
 import inspect
 import requests
 from requests import HTTPError, RequestException
+from requests.adapters import HTTPAdapter, Retry
 from typing import Any, Generator, List, Dict, TextIO, Optional, Union
 from .config import OnyxConfig
 from .field import OnyxField
@@ -14,24 +15,49 @@
 )
 from .endpoints import OnyxEndpoint
 
+RETRY_STATUS_CODES = [
+    502,  # Bad Gateway
+    503,  # Service Unavailable
+    504,  # Gateway Timeout
+]
+RETRY_METHODS = ["HEAD", "GET", "OPTIONS", "POST", "PUT", "PATCH", "DELETE"]
+
 
 class OnyxClientBase:
     __slots__ = "config", "_request_handler", "_session"
 
     def __init__(self, config: OnyxConfig):
         self.config = config
         self._session = None
-        self._request_handler = requests.request
+        self._request_handler = self._default_request_handler
 
     def __enter__(self):
-        self._session = requests.Session()
+        self._session = self._get_session()
         self._request_handler = self._session.request
         return self
 
     def __exit__(self, type, value, traceback):
         if self._session:
             self._session.close()
-        self._request_handler = requests.request
+        self._session = None
+        self._request_handler = self._default_request_handler
+
+    def _get_session(self) -> requests.Session:
+        session = requests.Session()
+        retry_strategy = Retry(
+            total=5,
+            backoff_factor=1,
+            status_forcelist=RETRY_STATUS_CODES,
+            allowed_methods=RETRY_METHODS,
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        session.mount("https://", adapter)
+        session.mount("http://", adapter)
+        return session
+
+    def _default_request_handler(self, method: str, **kwargs) -> requests.Response:
+        with self._get_session() as session:
+            return session.request(method, **kwargs)
 
     def _request(self, method: str, retries: int = 3, **kwargs) -> requests.Response:
         if not retries:

pyproject.toml

@@ -28,8 +28,8 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-dev = ["ruff", "pytest", "coverage"]
-test = ["pytest", "coverage"]
+dev = ["ruff", "pytest", "coverage", "responses"]
+test = ["pytest", "coverage", "responses"]
 
 [project.scripts]
 onyx = "onyx.cli:main"

tests/test_api.py

@@ -1,8 +1,11 @@
 import io
 import requests
+import responses
 import pytest
 from unittest import TestCase, mock
+from requests.adapters import HTTPAdapter
 from onyx import OnyxConfig, OnyxClient, exceptions, OnyxField, OnyxEndpoint
+from onyx.core import RETRY_STATUS_CODES
 
 
 DOMAIN = "https://onyx.domain"
@@ -1011,7 +1014,110 @@ def test_context_manager(self, mock_request):
                 client._session.request,  #  type: ignore
             )
 
-        self.assertEqual(client._request_handler, requests.request)
+        self.assertEqual(client._session, None)
+        self.assertEqual(client._request_handler, client._default_request_handler)
+
+    def test_retry_configured(self):
+        """
+        Test that the retry adapter is configured correctly.
+        """
+
+        def test_config(session):
+            # Verify the session has retry adapters configured for both http and https
+            for scheme in ("https://", "http://"):
+                adapter = session.get_adapter(scheme)
+                assert isinstance(adapter, HTTPAdapter)
+                self.assertEqual(adapter.max_retries.total, 5)
+                self.assertEqual(adapter.max_retries.backoff_factor, 1)
+
+                for status_code in RETRY_STATUS_CODES:
+                    self.assertIn(status_code, adapter.max_retries.status_forcelist)
+
+        # Test within context manager
+        with OnyxClient(self.config) as client:
+            test_config(client._session)
+
+        # Test outside context manager
+        # Capture the session being created
+        sessions = []
+
+        def _get_session():
+            session = self.client._get_session()
+            sessions.append(session)
+            return session
+
+        client = OnyxClient(self.config)
+        client._get_session = _get_session
+        with mock.patch.object(
+            requests.Session, "request", return_value=MockResponse(PROJECT_DATA)
+        ):
+            client.projects()
+        self.assertEqual(len(sessions), 1)
+        test_config(sessions[0])
+
+    @responses.activate
+    def test_retry(self):
+        """
+        Test that retries occur on selected status codes.
+        """
+
+        url = OnyxEndpoint["projects"](DOMAIN)
+
+        # Test within context manager
+        for status_code in RETRY_STATUS_CODES:
+            # First two requests return badly, third returns 200
+            responses.add(responses.GET, url, status=status_code)
+            responses.add(responses.GET, url, status=status_code)
+            responses.add(responses.GET, url, json=PROJECT_DATA, status=200)
+
+            with OnyxClient(self.config) as client:
+                result = client.projects()
+            self.assertEqual(result, PROJECT_DATA["data"])
+
+        # Test outside context manager
+        for status_code in RETRY_STATUS_CODES:
+            # First two requests return badly, third returns 200
+            responses.add(responses.GET, url, status=status_code)
+            responses.add(responses.GET, url, status=status_code)
+            responses.add(responses.GET, url, json=PROJECT_DATA, status=200)
+
+            result = OnyxClient(self.config).projects()
+            self.assertEqual(result, PROJECT_DATA["data"])
+
+        # 3 requests per status code * number of status codes * 2 (with/without context manager)
+        num_expected_calls = 3 * len(RETRY_STATUS_CODES) * 2
+        self.assertEqual(len(responses.calls), num_expected_calls)
+
+    @responses.activate
+    def test_retry_exhausted_error(self):
+        """
+        Test that an error is raised when all retries are exhausted.
+        """
+
+        url = OnyxEndpoint["projects"](DOMAIN)
+
+        # Test within context manager
+        for status_code in RETRY_STATUS_CODES:
+            # Return badly for all attempts (more than max retries)
+            for _ in range(10):
+                responses.add(responses.GET, url, status=status_code)
+
+            with OnyxClient(self.config) as client:
+                with pytest.raises(exceptions.OnyxConnectionError):
+                    client.projects()
+
+        # Test outside context manager
+        for status_code in RETRY_STATUS_CODES:
+            # Return badly for all attempts (more than max retries)
+            for _ in range(10):
+                responses.add(responses.GET, url, status=status_code)
+
+            with pytest.raises(exceptions.OnyxConnectionError):
+                OnyxClient(self.config).projects()
+
+        # (1 initial request per test + 5 retries) * number of status codes * 2 (with/without context manager)
+        num_expected_calls = (1 + 5) * len(RETRY_STATUS_CODES) * 2
+        self.assertEqual(len(responses.calls), num_expected_calls)
 
     @mock.patch("onyx.OnyxClient._request_handler", side_effect=mock_request)
     def test_connection_error(self, mock_request):