diff --git a/matrix/api.py b/matrix/api.py index 43074bf..b078bf4 100644 --- a/matrix/api.py +++ b/matrix/api.py @@ -1,21 +1,23 @@ import asyncio -from typing import Awaitable, Iterable, TypeVar +from typing import Awaitable, Callable, Iterable, TypeVar from nio import ErrorResponse, Response -from matrix.errors import MatrixError +from matrix.errors import MatrixApiError T = TypeVar("T", bound=Response) R = TypeVar("R") async def matrix_call(coro: Awaitable[T], /, *, error_message: str) -> T: - """Await `coro`, translating any failure into a `MatrixError`. + """Await `coro`, translating any failure into a `MatrixApiError`. matrix-nio's `AsyncClient` methods don't raise on API-level errors; they return an `ErrorResponse` instead of raising. This wraps a single call so both transport-level exceptions and nio `ErrorResponse` results become a - `MatrixError` carrying `error_message`. + `MatrixApiError` carrying `error_message`. When the `ErrorResponse` is a + rate limit with a `retry_after_ms`, that value is carried through + unchanged on `MatrixApiError.retry_after_ms`. ## Example @@ -29,14 +31,54 @@ async def matrix_call(coro: Awaitable[T], /, *, error_message: str) -> T: try: response = await coro except Exception as e: - raise MatrixError(f"{error_message}: {e}") from e + raise MatrixApiError(f"{error_message}: {e}") from e if isinstance(response, ErrorResponse): - raise MatrixError(f"{error_message}: {response}") + raise MatrixApiError( + f"{error_message}: {response}", retry_after_ms=response.retry_after_ms + ) return response +async def with_retry( + func: Callable[[], Awaitable[R]], + /, + *, + retries: int = 3, + base_delay: float = 1.0, +) -> R: + """Call `func`, retrying on `MatrixApiError` with backoff. + + Retries up to `retries` times after the initial attempt (so up to + `retries + 1` total calls). If the raised `MatrixApiError` carries a + `retry_after_ms` (from the Matrix server), that wait is used exactly + (converted to seconds); otherwise falls back to exponential delay + (`base_delay * 2 ** attempt`). Only `MatrixApiError` triggers a retry; + any other exception propagates. + + ## Example + + ```python + message = await with_retry(lambda: room.send(content, notice=notice)) + ``` + """ + for attempt in range(retries + 1): + try: + return await func() + except MatrixApiError as e: + if attempt == retries: + raise + + delay = base_delay * 2**attempt + if e.retry_after_ms is not None: + delay = e.retry_after_ms / 1000 + + await asyncio.sleep(delay) + + raise AssertionError("unreachable") + + async def bounded_gather( coros: Iterable[Awaitable[R]], /, diff --git a/matrix/errors.py b/matrix/errors.py index c9cca60..f3f23cf 100644 --- a/matrix/errors.py +++ b/matrix/errors.py @@ -13,6 +13,12 @@ class MatrixError(Exception): pass +class MatrixApiError(MatrixError): + def __init__(self, message: str, *, retry_after_ms: int | None = None): + super().__init__(message) + self.retry_after_ms = retry_after_ms + + class RoomNotFoundError(MatrixError): pass diff --git a/tests/test_api.py b/tests/test_api.py index 1a2ad8c..9098919 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,11 +1,12 @@ import asyncio import inspect +from unittest.mock import AsyncMock, patch import pytest from nio import RoomSendResponse, RoomSendError -from matrix.errors import MatrixError -from matrix.api import bounded_gather, matrix_call +from matrix.errors import ConfigError, MatrixApiError, MatrixError +from matrix.api import bounded_gather, matrix_call, with_retry @pytest.mark.asyncio @@ -19,11 +20,11 @@ async def call(): @pytest.mark.asyncio -async def test_matrix_call_with_transport_exception__expect_matrix_error(): +async def test_matrix_call_with_transport_exception__expect_matrix_api_error(): async def call(): raise Exception("Network error") - with pytest.raises(MatrixError, match="Failed to send message: Network error"): + with pytest.raises(MatrixApiError, match="Failed to send message: Network error"): await matrix_call(call(), error_message="Failed to send message") @@ -36,6 +37,30 @@ async def call(): await matrix_call(call(), error_message="Failed to send message") +@pytest.mark.asyncio +async def test_matrix_call_with_rate_limited_error_response__expect_retry_after_ms_set(): + async def call(): + return RoomSendError( + "too many requests", "M_LIMIT_EXCEEDED", retry_after_ms=5000 + ) + + with pytest.raises(MatrixApiError) as exc_info: + await matrix_call(call(), error_message="Failed to send message") + + assert exc_info.value.retry_after_ms == 5000 + + +@pytest.mark.asyncio +async def test_matrix_call_with_error_response_without_retry_after__expect_retry_after_ms_none(): + async def call(): + return RoomSendError("not allowed", "M_FORBIDDEN") + + with pytest.raises(MatrixApiError) as exc_info: + await matrix_call(call(), error_message="Failed to send message") + + assert exc_info.value.retry_after_ms is None + + def test_matrix_call_requires_keyword_error_message__expect_type_error(): with pytest.raises(TypeError): matrix_call(None, "Failed to send message") @@ -46,6 +71,103 @@ def test_matrix_call_requires_positional_coro__expect_type_error(): matrix_call(coro=None, error_message="Failed to send message") +@pytest.mark.asyncio +async def test_with_retry_with_success_first_try__expect_response_returned(): + call = AsyncMock(return_value="ok") + + result = await with_retry(lambda: call()) + + assert result == "ok" + assert call.await_count == 1 + + +@pytest.mark.asyncio +async def test_with_retry_with_transient_matrix_error_then_success__expect_response_returned_after_retry(): + call = AsyncMock(side_effect=[MatrixApiError("rate limited"), "ok"]) + + with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep: + result = await with_retry(lambda: call()) + + assert result == "ok" + assert call.await_count == 2 + sleep.assert_awaited_once_with(1.0) + + +@pytest.mark.asyncio +async def test_with_retry_exhausts_retries__expect_matrix_api_error_raised(): + call = AsyncMock(side_effect=MatrixApiError("still failing")) + + with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(MatrixApiError, match="still failing"): + await with_retry(lambda: call(), retries=2) + + assert call.await_count == 3 + + +@pytest.mark.asyncio +async def test_with_retry_with_non_matrix_exception__expect_immediate_raise_no_retry(): + call = AsyncMock(side_effect=Exception("boom")) + + with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep: + with pytest.raises(Exception, match="boom"): + await with_retry(lambda: call()) + + assert call.await_count == 1 + sleep.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_with_retry_with_non_api_matrix_error__expect_immediate_raise_no_retry(): + call = AsyncMock(side_effect=ConfigError("token")) + + with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep: + with pytest.raises(ConfigError): + await with_retry(lambda: call()) + + assert call.await_count == 1 + sleep.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_with_retry_uses_exponential_backoff__expect_increasing_delays(): + call = AsyncMock( + side_effect=[ + MatrixApiError("fail 1"), + MatrixApiError("fail 2"), + MatrixApiError("fail 3"), + "ok", + ] + ) + + with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep: + result = await with_retry(lambda: call(), retries=3, base_delay=1.0) + + assert result == "ok" + assert [call_args.args[0] for call_args in sleep.await_args_list] == [ + 1.0, + 2.0, + 4.0, + ] + + +@pytest.mark.asyncio +async def test_with_retry_with_retry_after_ms__expect_server_delay_honored(): + call = AsyncMock( + side_effect=[MatrixApiError("rate limited", retry_after_ms=2500), "ok"] + ) + + with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep: + result = await with_retry(lambda: call()) + + assert result == "ok" + sleep.assert_awaited_once_with(2.5) + + +def test_with_retry_requires_positional_func__expect_type_error(): + with pytest.raises(TypeError): + with_retry(func=lambda: None) + + @pytest.mark.asyncio async def test_bounded_gather_with_success__expect_results_in_order(): async def value(n):