Add MessageEvent.await_with_auto_ack() (#8)

Author: vrslev

README.md

@@ -81,30 +81,32 @@ async for event in client.listen_to_events():
             print(f"{short_description}:\n{body!s}")
 ```
 
-More complex example, that involves handling all possible events:
+More complex example, that involves handling all possible events, and auto-acknowledgement:
 
 ```python
 async with asyncio.TaskGroup() as task_group:
     async for event in client.listen_to_events():
         match event:
             case stompman.MessageEvent(body=body):
-                # Validate message ASAP and ack/nack, so that server won't assume we're not reliable
-                try:
-                    validated_message = MyMessageModel.model_validate_json(body)
-                except ValidationError:
-                    await event.nack()
-                    raise
-
-                await event.ack()
-                task_group.create_task(run_business_logic(validated_message))
+                task_group.create_task(event.await_with_auto_ack(handle_message(body)))
             case stompman.ErrorEvent(message_header=short_description, body=body):
                 logger.error(
                     "Received an error from server", short_description=short_description, body=body, event=event
                 )
             case stompman.HeartbeatEvent():
                 task_group.create_task(update_healthcheck_status())
+
+
+async def handle_message(body: bytes) -> None:
+    try:
+        validated_message = MyMessageModel.model_validate_json(body)
+        await run_business_logic(validated_message)
+    except Exception:
+        logger.exception("Failed to handle message", body=body)
 ```
 
+You can pass awaitable object (coroutine, for example) to `Message.await_with_auto_ack()`. In case of error, it will catch any exceptions, send NACK to server and propagate them to the caller. Otherwise, it will send ACK, acknowledging the message was processed successfully.
+
 ### Cleaning Up
 
 stompman takes care of cleaning up resources automatically. When you leave the context of async context managers `stompman.Client()`, `client.subscribe()`, or `client.enter_transaction()`, the necessary frames will be sent to the server.

stompman/listening_events.py

@@ -1,3 +1,4 @@
+from collections.abc import Awaitable
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
@@ -36,6 +37,21 @@ async def nack(self) -> None:
             )
         )
 
+    async def await_with_auto_ack(
+        self, awaitable: Awaitable[None], exception_types: tuple[type[Exception],] = (Exception,)
+    ) -> None:
+        called_nack = False
+
+        try:
+            await awaitable
+        except exception_types:
+            await self.nack()
+            called_nack = True
+            raise
+        finally:
+            if not called_nack:
+                await self.ack()
+
 
 @dataclass
 class ErrorEvent:

tests/test_client.py

@@ -346,6 +346,59 @@ async def test_ack_nack() -> None:
     assert_frames_between_lifespan_match(collected_frames, [message_frame, nack_frame, ack_frame])
 
 
+def get_mocked_message_event() -> tuple[MessageEvent, mock.AsyncMock, mock.AsyncMock]:
+    ack_mock, nack_mock = mock.AsyncMock(), mock.AsyncMock()
+
+    class CustomMessageEvent(MessageEvent):
+        ack = ack_mock
+        nack = nack_mock
+
+    return (
+        CustomMessageEvent(
+            _frame=MessageFrame(
+                headers={"destination": "destination", "message-id": "message-id", "subscription": "subscription"},
+                body=b"",
+            ),
+            _client=mock.Mock(),
+        ),
+        ack_mock,
+        nack_mock,
+    )
+
+
+async def test_message_event_await_with_auto_ack_nack() -> None:
+    event, ack, nack = get_mocked_message_event()
+
+    async def raises_runtime_error() -> None:  # noqa: RUF029
+        raise RuntimeError
+
+    with suppress(RuntimeError):
+        await event.await_with_auto_ack(raises_runtime_error(), exception_types=(Exception,))
+
+    ack.assert_not_called()
+    nack.assert_called_once_with()
+
+
+async def test_message_event_await_with_auto_ack_ack_raises() -> None:
+    event, ack, nack = get_mocked_message_event()
+
+    async def func() -> None:  # noqa: RUF029
+        raise Exception  # noqa: TRY002
+
+    with suppress(Exception):
+        await event.await_with_auto_ack(func(), exception_types=(RuntimeError,))
+
+    ack.assert_called_once_with()
+    nack.assert_not_called()
+
+
+async def test_message_event_await_with_auto_ack_ack_ok() -> None:
+    event, ack, nack = get_mocked_message_event()
+    await event.await_with_auto_ack(mock.AsyncMock()())
+    ack.assert_called_once_with()
+    nack.assert_not_called()
+
+
 async def test_send_message_and_enter_transaction_ok(monkeypatch: pytest.MonkeyPatch) -> None:
     body = b"hello"
     destination = "/queue/test"