From d26be36a585eca730e0744bc99d5ef3856d48edc Mon Sep 17 00:00:00 2001 From: Alexandra Szanto Date: Tue, 3 Nov 2026 15:21:12 +0100 Subject: [PATCH] refactor: extract aiohttp transport --- UnleashClient/_aiohttp_transport.py | 71 +++++++ UnleashClient/_async_unleash_api_client.py | 185 +++++++++--------- UnleashClient/_transport.py | 13 ++ tests/unit_tests/test_aiohttp_transport.py | 179 +++++++++++++++++ tests/unit_tests/test_async_metrics.py | 2 +- .../test_async_unleash_api_client.py | 95 ++++++++- 6 files changed, 446 insertions(+), 99 deletions(-) create mode 100644 UnleashClient/_aiohttp_transport.py create mode 100644 tests/unit_tests/test_aiohttp_transport.py diff --git a/UnleashClient/_aiohttp_transport.py b/UnleashClient/_aiohttp_transport.py new file mode 100644 index 00000000..8ae90da7 --- /dev/null +++ b/UnleashClient/_aiohttp_transport.py @@ -0,0 +1,71 @@ +""" +HTTP layer under the async Unleash API client: sends one request with aiohttp +and returns the response. +""" + +from typing import Optional + +from UnleashClient._transport import Request, Response +from UnleashClient.constants import REQUEST_TIMEOUT +from UnleashClient.errors import AlreadyClosedError + +try: + import aiohttp +except ImportError as exc: # pragma: no cover + raise ImportError( + "AsyncUnleashClient requires aiohttp, which is not installed. " + "Install it with: pip install UnleashClient[async]" + ) from exc + + +class AiohttpTransport: + """ + Sends requests with ``aiohttp``, on one session reused across requests. + + The session is opened on the first request, since aiohttp needs a running + event loop to create one, and closed by :meth:`aclose`. A closed transport + stays closed. + + :param timeout: Seconds to wait for the connection and for each read. + """ + + def __init__(self, *, timeout: float = REQUEST_TIMEOUT) -> None: + self._timeout = timeout + self._session: Optional["aiohttp.ClientSession"] = None + self._is_closed: bool = False + + def _get_session(self) -> "aiohttp.ClientSession": + if self._is_closed: + raise AlreadyClosedError() + if self._session is None: + self._session = aiohttp.ClientSession() + return self._session + + def _client_timeout(self) -> "aiohttp.ClientTimeout": + # sock_connect/sock_read rather than total: the sync client's scalar + # timeout bounds each socket operation, not the whole request. + return aiohttp.ClientTimeout( + total=None, sock_connect=self._timeout, sock_read=self._timeout + ) + + async def execute(self, request: Request) -> Response: + async with self._get_session().request( + request.method, + request.url, + headers=dict(request.headers), + params=dict(request.params) or None, + data=request.body, + timeout=self._client_timeout(), + ) as resp: + body = await resp.read() + return Response(status=resp.status, headers=resp.headers, body=body) + + async def aclose(self) -> None: + """ + Close the session. Every later request raises + :class:`AlreadyClosedError`. Closing a closed transport does nothing. + """ + self._is_closed = True + session, self._session = self._session, None + if session is not None: + await session.close() diff --git a/UnleashClient/_async_unleash_api_client.py b/UnleashClient/_async_unleash_api_client.py index c2885253..2b34961d 100644 --- a/UnleashClient/_async_unleash_api_client.py +++ b/UnleashClient/_async_unleash_api_client.py @@ -1,18 +1,22 @@ -"""Unleash API client over aiohttp, for the async Unleash client.""" +"""Unleash API client for the async Unleash client.""" import asyncio import json import threading from typing import Any, Dict, Optional, Tuple +from UnleashClient._aiohttp_transport import AiohttpTransport from UnleashClient._event_dispatcher import _EventDispatcher from UnleashClient._headers import _HeaderFactory +from UnleashClient._transport import AsyncTransport, Request from UnleashClient._unleash_api_client import ( FetchResult, _emit_fetch_completed, _failure_status, _fetched_status, + _log_resp_info, _normalized_url, + _text, ) from UnleashClient.config import PollingOptions, UnleashConfig from UnleashClient.constants import FEATURES_URL, METRICS_URL, REGISTER_URL @@ -28,7 +32,6 @@ "Install it with: pip install UnleashClient[async]" ) from exc - RETRY_STATUSES = frozenset({500, 502, 504}) RETRYABLE_ERRORS = (aiohttp.ClientConnectionError, asyncio.TimeoutError) @@ -36,12 +39,6 @@ FATAL_URL_ERRORS = (aiohttp.InvalidURL, aiohttp.NonHttpUrlClientError) -async def _log_resp_info(resp: "aiohttp.ClientResponse") -> None: - LOGGER.debug("HTTP status code: %s", resp.status) - LOGGER.debug("HTTP headers: %s", resp.headers) - LOGGER.debug("HTTP content: %s", await resp.text()) - - class _AsyncUnleashApiClient: """The asyncio twin of :class:`UnleashClient._unleash_api_client._UnleashApiClient`.""" @@ -50,16 +47,27 @@ def __init__( config: UnleashConfig, headers: _HeaderFactory, events: Optional[_EventDispatcher] = None, + transport: Optional[AsyncTransport] = None, ) -> None: """ :param config: read for the url, timeouts, retries and project. :param headers: builds the header set each request needs. :param events: receives ``FETCH_COMPLETED`` after every feature fetch. + :param transport: sends the requests. Defaults to an + :class:`AiohttpTransport` built from ``config.http``, + which :meth:`aclose` closes. A transport passed in is + left open: whoever created it closes it. """ self._config: UnleashConfig = config self._headers: _HeaderFactory = headers self._events: Optional[_EventDispatcher] = events - self._session: Optional["aiohttp.ClientSession"] = None + # set only when this API client built the transport, and so it closes it. + self._default_transport: Optional[AiohttpTransport] = None + if transport is None: + transport = self._default_transport = AiohttpTransport( + timeout=config.http.timeout + ) + self._transport: AsyncTransport = transport self._is_closed: bool = False self._close_lock = threading.Lock() @@ -67,22 +75,10 @@ def _raise_if_closed(self) -> None: if self._is_closed: raise AlreadyClosedError() - async def _get_session(self) -> "aiohttp.ClientSession": - if self._session is None or self._session.closed: - self._session = aiohttp.ClientSession() - return self._session - - def _timeout(self) -> "aiohttp.ClientTimeout": - # sock_connect/sock_read rather than total: the sync client's scalar - # timeout bounds each socket operation, not the whole request. - seconds = self._config.http.timeout - return aiohttp.ClientTimeout( - total=None, sock_connect=seconds, sock_read=seconds - ) - async def aclose(self) -> None: """ - Close the pooled session and put the API client out of service. + Put the API client out of service, and close its transport if the API + client created it. Every later request raises :class:`AlreadyClosedError`. Closing an already closed API client does nothing. @@ -92,10 +88,9 @@ async def aclose(self) -> None: return self._is_closed = True - session, self._session = self._session, None - if session is not None and not session.closed: - await session.close() + if self._default_transport is not None: + await self._default_transport.aclose() async def fetch_features( self, etag: str = "", *, options: PollingOptions @@ -140,51 +135,47 @@ async def _fetch_features( if options.project_name: base_params = {"project": options.project_name} - session = await self._get_session() + request = Request( + method="GET", url=base_url, headers=headers, params=base_params + ) attempts = 1 + max(config.http.retries, 0) for attempt in range(attempts): last_attempt = attempt == attempts - 1 try: - async with session.get( - base_url, - headers=headers, - params=base_params, - timeout=self._timeout(), - ) as resp: - if resp.status in RETRY_STATUSES and not last_attempt: - continue - - http_status_code = resp.status - if resp.status not in [200, 304]: - await _log_resp_info(resp) - LOGGER.warning( - "Unleash Client feature fetch failed due to unexpected HTTP status code: %s", - resp.status, - ) - raise Exception( # pylint: disable=broad-exception-raised - "Unleash Client feature fetch failed!" - ) - - fetched_etag = resp.headers.get("etag", "") - - if resp.status == 304: - return ( - FetchResult(None, fetched_etag, not_modified=True), - FetchStatus.NOT_CHANGED, - ) - - raw_state = await resp.text() - return ( - FetchResult(raw_state, fetched_etag), - _fetched_status(raw_state), - ) + resp = await self._transport.execute(request) except RETRYABLE_ERRORS: - # No usable response arrived, even if a status line did - # before the body timed out - http_status_code = None if last_attempt: raise + continue + + if resp.status in RETRY_STATUSES and not last_attempt: + continue + + http_status_code = resp.status + if resp.status not in [200, 304]: + _log_resp_info(resp) + LOGGER.warning( + "Unleash Client feature fetch failed due to unexpected HTTP status code: %s", + resp.status, + ) + raise Exception( # pylint: disable=broad-exception-raised + "Unleash Client feature fetch failed!" + ) + + fetched_etag = resp.headers.get("etag", "") + + if resp.status == 304: + return ( + FetchResult(None, fetched_etag, not_modified=True), + FetchStatus.NOT_CHANGED, + ) + + raw_state = _text(resp) + return ( + FetchResult(raw_state, fetched_etag), + _fetched_status(raw_state), + ) except Exception as exc: LOGGER.exception( "Unleash Client feature fetch failed due to exception: %s", exc @@ -211,24 +202,26 @@ async def register(self, payload: Dict[str, Any]) -> bool: LOGGER.info("Registering unleash client with unleash @ %s", config.url) LOGGER.info("Registration request information: %s", payload) - session = await self._get_session() - async with session.post( - _normalized_url(config.url, REGISTER_URL), - data=json.dumps(payload), - headers=self._headers.base(), - timeout=self._timeout(), - ) as resp: - if resp.status not in {200, 202}: - await _log_resp_info(resp) - LOGGER.warning( - "Unleash Client registration failed due to unexpected HTTP status code: %s", - resp.status, - ) - return False + resp = await self._transport.execute( + Request( + method="POST", + url=_normalized_url(config.url, REGISTER_URL), + headers=self._headers.base(), + body=json.dumps(payload).encode("utf-8"), + ) + ) - LOGGER.info("Unleash Client successfully registered!") + if resp.status not in {200, 202}: + _log_resp_info(resp) + LOGGER.warning( + "Unleash Client registration failed due to unexpected HTTP status code: %s", + resp.status, + ) + return False - return True + LOGGER.info("Unleash Client successfully registered!") + + return True except FATAL_URL_ERRORS as exc: LOGGER.exception( "Unleash Client registration failed fatally due to exception: %s", exc @@ -259,24 +252,26 @@ async def send_metrics(self, payload: Dict[str, Any]) -> bool: LOGGER.info("Sending messages to with unleash @ %s", config.url) LOGGER.info("unleash metrics information: %s", payload) - session = await self._get_session() - async with session.post( - _normalized_url(config.url, METRICS_URL), - data=json.dumps(payload), - headers=self._headers.metrics(), - timeout=self._timeout(), - ) as resp: - if resp.status != 202: - await _log_resp_info(resp) - LOGGER.warning( - "Unleash Client metrics submission due to unexpected HTTP status code: %s", - resp.status, - ) - return False + resp = await self._transport.execute( + Request( + method="POST", + url=_normalized_url(config.url, METRICS_URL), + headers=self._headers.metrics(), + body=json.dumps(payload).encode("utf-8"), + ) + ) + + if resp.status != 202: + _log_resp_info(resp) + LOGGER.warning( + "Unleash Client metrics submission due to unexpected HTTP status code: %s", + resp.status, + ) + return False - LOGGER.info("Unleash Client metrics successfully sent!") + LOGGER.info("Unleash Client metrics successfully sent!") - return True + return True except aiohttp.ClientError as exc: LOGGER.warning( "Unleash Client metrics submission failed due to exception: %s", exc diff --git a/UnleashClient/_transport.py b/UnleashClient/_transport.py index 6bc9c2ce..ac647819 100644 --- a/UnleashClient/_transport.py +++ b/UnleashClient/_transport.py @@ -61,6 +61,19 @@ def execute(self, request: Request) -> Response: """ +@runtime_checkable +class AsyncTransport(Protocol): + """Sends requests for the async client.""" + + async def execute(self, request: Request) -> Response: + """ + Send ``request`` once and return the response, whatever its status. + + :raises Exception: only when no response arrived, e.g. a connection + error or a timeout. + """ + + class RequestsTransport: """ Sends requests with ``requests``, on a fresh session for every request. diff --git a/tests/unit_tests/test_aiohttp_transport.py b/tests/unit_tests/test_aiohttp_transport.py new file mode 100644 index 00000000..ea02b01d --- /dev/null +++ b/tests/unit_tests/test_aiohttp_transport.py @@ -0,0 +1,179 @@ +import asyncio +import json + +import aiohttp +import pytest +import pytest_asyncio +from pytest import mark + +from tests.utilities.fake_unleash_server import FakeUnleash +from UnleashClient._aiohttp_transport import AiohttpTransport +from UnleashClient._transport import AsyncTransport, Request +from UnleashClient.errors import AlreadyClosedError + +PATH = "/client/features" + + +@pytest_asyncio.fixture +async def server(): + fake = FakeUnleash() + await fake.start() + try: + yield fake + finally: + await fake.close() + + +@pytest_asyncio.fixture +async def transport(): + transport = AiohttpTransport() + try: + yield transport + finally: + await transport.aclose() + + +def test_aiohttp_transport_is_an_async_transport(): + assert isinstance(AiohttpTransport(), AsyncTransport) + + +@mark.asyncio +async def test_execute_sends_a_get_with_its_headers_and_params(server, transport): + server.on("GET", PATH) + + await transport.execute( + Request( + method="GET", + url=server.base_url + PATH, + headers={"X-Test": "value"}, + params={"project": "ivan"}, + ) + ) + + sent = server.calls("GET", PATH)[0] + assert sent.query == {"project": "ivan"} + assert sent.headers["X-Test"] == "value" + assert sent.body == "" + + +@mark.asyncio +async def test_execute_sends_a_post_with_its_body(server, transport): + server.on("POST", PATH, status=202) + + await transport.execute( + Request( + method="POST", + url=server.base_url + PATH, + headers={}, + body=b'{"appName": "x"}', + ) + ) + + assert json.loads(server.calls("POST", PATH)[0].body) == {"appName": "x"} + + +@mark.asyncio +async def test_execute_returns_the_status_headers_and_body(server, transport): + server.on("GET", PATH, payload={"version": 1}, headers={"ETag": "v1"}) + + response = await transport.execute( + Request(method="GET", url=server.base_url + PATH, headers={}) + ) + + assert response.status == 200 + assert json.loads(response.body) == {"version": 1} + assert response.headers["etag"] == "v1" + + +@mark.asyncio +@mark.parametrize("status", [304, 401, 500]) +async def test_execute_returns_any_status_instead_of_raising(server, transport, status): + server.on("GET", PATH, status=status) + + response = await transport.execute( + Request(method="GET", url=server.base_url + PATH, headers={}) + ) + + assert response.status == status + + +@mark.asyncio +async def test_execute_raises_when_no_response_arrives(server, transport): + server.on("GET", PATH, disconnect=True, repeat=True) + + with pytest.raises(aiohttp.ClientConnectionError): + await transport.execute( + Request(method="GET", url=server.base_url + PATH, headers={}) + ) + + +@mark.asyncio +async def test_execute_raises_when_the_server_does_not_answer_in_time(server): + server.on("GET", PATH, hang=True) + transport = AiohttpTransport(timeout=0.1) + + try: + with pytest.raises(asyncio.TimeoutError): + await transport.execute( + Request(method="GET", url=server.base_url + PATH, headers={}) + ) + finally: + await transport.aclose() + + +@mark.asyncio +async def test_execute_reuses_one_session(server, transport): + server.on("GET", PATH, repeat=True) + request = Request(method="GET", url=server.base_url + PATH, headers={}) + + await transport.execute(request) + first = transport._session + await transport.execute(request) + + assert first is not None + assert transport._session is first + + +@mark.asyncio +async def test_aclose_closes_the_session(server, transport): + server.on("GET", PATH) + await transport.execute( + Request(method="GET", url=server.base_url + PATH, headers={}) + ) + session = transport._session + + await transport.aclose() + + assert session.closed + + +@mark.asyncio +async def test_execute_after_aclose_raises(server, transport): + server.on("GET", PATH, repeat=True) + request = Request(method="GET", url=server.base_url + PATH, headers={}) + await transport.execute(request) + + await transport.aclose() + + with pytest.raises(AlreadyClosedError): + await transport.execute(request) + assert len(server.calls("GET", PATH)) == 1 + + +@mark.asyncio +async def test_execute_after_aclose_raises_without_an_earlier_request(server): + transport = AiohttpTransport() + + await transport.aclose() + + with pytest.raises(AlreadyClosedError): + await transport.execute( + Request(method="GET", url=server.base_url + PATH, headers={}) + ) + assert server.requests == [] + + +@mark.asyncio +async def test_closing_a_closed_transport_does_nothing(transport): + await transport.aclose() + await transport.aclose() diff --git a/tests/unit_tests/test_async_metrics.py b/tests/unit_tests/test_async_metrics.py index 9b05823f..ab174c87 100644 --- a/tests/unit_tests/test_async_metrics.py +++ b/tests/unit_tests/test_async_metrics.py @@ -309,7 +309,7 @@ async def test_constructing_the_reporter_registers_no_job_and_opens_no_session( reporter = build_reporter() assert reporter._scheduler.registered == [] - assert reporter._api._session is None + assert reporter._api._transport._session is None # impact metrics diff --git a/tests/unit_tests/test_async_unleash_api_client.py b/tests/unit_tests/test_async_unleash_api_client.py index 6c9a873a..5d5be8bc 100644 --- a/tests/unit_tests/test_async_unleash_api_client.py +++ b/tests/unit_tests/test_async_unleash_api_client.py @@ -21,9 +21,11 @@ PROJECT_NAME, REQUEST_RETRIES, REQUEST_TIMEOUT, + URL, ) from UnleashClient._async_unleash_api_client import _AsyncUnleashApiClient from UnleashClient._headers import _HeaderFactory +from UnleashClient._transport import Request, Response from UnleashClient.config import ( HttpOptions, IdentityOptions, @@ -377,11 +379,11 @@ async def test_the_session_is_reused_across_requests(server, api_client): server.on("POST", METRICS_PATH, status=202, payload={}, repeat=True) await api_client.send_metrics(MOCK_METRICS_REQUEST) - first = api_client._session + first = api_client._transport._session await api_client.send_metrics(MOCK_METRICS_REQUEST) assert first is not None - assert api_client._session is first + assert api_client._transport._session is first @mark.asyncio @@ -483,9 +485,96 @@ async def test_fetch_features_reports_unavailable_when_the_body_times_out( async def body_times_out(_response): raise asyncio.TimeoutError() - monkeypatch.setattr(aiohttp.ClientResponse, "text", body_times_out) + monkeypatch.setattr(aiohttp.ClientResponse, "read", body_times_out) events = RecordingEvents() await build_api_client(events=events).fetch_features(options=PollingOptions()) assert fetch_outcomes(events) == [FetchStatus.UNAVAILABLE] + + +class RecordingTransport: + """Records every request and answers each with the same response.""" + + def __init__(self, response: Response): + self.response = response + self.requests = [] + self.closed = False + + async def execute(self, request: Request) -> Response: + self.requests.append(request) + return self.response + + async def aclose(self) -> None: + self.closed = True + + +def api_client_over(transport: RecordingTransport) -> _AsyncUnleashApiClient: + config = UnleashConfig( + url=URL, + app_name=APP_NAME, + identity=IdentityOptions(instance_id=INSTANCE_ID), + ) + return _AsyncUnleashApiClient(config, _HeaderFactory(config), transport=transport) + + +@mark.asyncio +async def test_fetch_features_sends_its_request_through_the_transport(): + transport = RecordingTransport( + Response(status=200, headers={"etag": ETAG_VALUE}, body=b'{"version": 1}') + ) + + result = await api_client_over(transport).fetch_features( + ETAG_VALUE, options=PollingOptions(project_name=PROJECT_NAME) + ) + + assert result == ('{"version": 1}', ETAG_VALUE, False) + (request,) = transport.requests + assert request.method == "GET" + assert request.url == URL + FEATURES_URL + assert request.params == {"project": PROJECT_NAME} + assert request.headers["If-None-Match"] == ETAG_VALUE + + +@mark.asyncio +async def test_register_posts_its_payload_through_the_transport(): + transport = RecordingTransport(Response(status=202, headers={}, body=b"")) + + assert await api_client_over(transport).register({"appName": APP_NAME}) + + (request,) = transport.requests + assert request.method == "POST" + assert request.url == URL + REGISTER_URL + assert json.loads(request.body) == {"appName": APP_NAME} + + +@mark.asyncio +async def test_send_metrics_posts_its_payload_through_the_transport(): + transport = RecordingTransport(Response(status=202, headers={}, body=b"")) + + assert await api_client_over(transport).send_metrics(MOCK_METRICS_REQUEST) + + (request,) = transport.requests + assert request.method == "POST" + assert request.url == URL + METRICS_URL + assert json.loads(request.body) == MOCK_METRICS_REQUEST + + +@mark.asyncio +async def test_aclose_leaves_a_transport_it_was_given_open(): + transport = RecordingTransport(Response(status=202, headers={}, body=b"")) + + await api_client_over(transport).aclose() + + assert not transport.closed + + +@mark.asyncio +async def test_aclose_closes_the_default_transport(server, api_client): + server.on("POST", METRICS_PATH, status=202, payload={}) + await api_client.send_metrics(MOCK_METRICS_REQUEST) + session = api_client._transport._session + + await api_client.aclose() + + assert session.closed