diff --git a/onyx/core.py b/onyx/core.py index b0e5271..3718a76 100644 --- a/onyx/core.py +++ b/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,6 +15,13 @@ ) 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" @@ -21,17 +29,36 @@ class OnyxClientBase: 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, + raise_on_status=False, + ) + 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: diff --git a/pyproject.toml b/pyproject.toml index 2f1c861..273b4a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "climb-onyx-client" -version = "4.9.0" +version = "4.9.1" description = "CLI and Python library for Onyx" readme = "README.md" requires-python = ">=3.9" @@ -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" diff --git a/tests/test_api.py b/tests/test_api.py index e8f73b6..b7be1ed 100644 --- a/tests/test_api.py +++ b/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.OnyxServerError): + 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.OnyxServerError): + 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):