Added additional parameters to HandlerException

faststream/broker/acknowledgement_watcher.py

@@ -151,16 +151,13 @@ async def __aexit__(
                 self.watcher.remove(self.message.message_id)
 
             elif isinstance(exc_val, AckMessage):
-                await self.__ack()
+                await self.__ack(**exc_val.extra_options)
 
             elif isinstance(exc_val, NackMessage):
-                if self.watcher.is_max(self.message.message_id):
-                    await self.__reject()
-                else:
-                    await self.__nack()
+                await self.__nack(**exc_val.extra_options)
 
             elif isinstance(exc_val, RejectMessage):  # pragma: no branch
-                await self.__reject()
+                await self.__reject(**exc_val.extra_options)
 
             # Exception was processed and suppressed
             return True
@@ -174,25 +171,25 @@ async def __aexit__(
         # Exception was not processed
         return False
 
-    async def __ack(self) -> None:
+    async def __ack(self, **exc_extra_options: Any) -> None:
         try:
-            await self.message.ack(**self.extra_options)
+            await self.message.ack(**self.extra_options, **exc_extra_options)
         except Exception as er:
             if self.logger is not None:
                 self.logger.log(logging.ERROR, er, exc_info=er)
         else:
             self.watcher.remove(self.message.message_id)
 
-    async def __nack(self) -> None:
+    async def __nack(self, **exc_extra_options: Any) -> None:
         try:
-            await self.message.nack(**self.extra_options)
+            await self.message.nack(**self.extra_options, **exc_extra_options)
         except Exception as er:
             if self.logger is not None:
                 self.logger.log(logging.ERROR, er, exc_info=er)
 
-    async def __reject(self) -> None:
+    async def __reject(self, **exc_extra_options: Any) -> None:
         try:
-            await self.message.reject(**self.extra_options)
+            await self.message.reject(**self.extra_options, **exc_extra_options)
         except Exception as er:
             if self.logger is not None:
                 self.logger.log(logging.ERROR, er, exc_info=er)

faststream/exceptions.py

@@ -1,4 +1,4 @@
-from typing import Iterable
+from typing import Any, Iterable
 
 
 class FastStreamException(Exception):  # noqa: N818
@@ -35,21 +35,57 @@ def __str__(self) -> str:
 
 
 class AckMessage(HandlerException):
-    """Raise it to `ack` a message immediately."""
+    """Exception raised to acknowledge a message immediately.
+
+    This exception can be used to ack a message with additional options.
+    To watch all allowed parameters, please take a look at your broker `message.ack(**extra_options)` method
+    signature.
+
+    Args:
+        extra_options (Any): Additional parameters that will be passed to `message.ack(**extra_options)` method.
+    """
+
+    def __init__(self, **extra_options: Any):
+        self.extra_options = extra_options
+        super().__init__()
 
     def __str__(self) -> str:
         return "Message was acked"
 
 
 class NackMessage(HandlerException):
-    """Raise it to `nack` a message immediately."""
+    """Exception raised to negatively acknowledge a message immediately.
+
+    This exception can be used to nack a message with additional options.
+    To watch all allowed parameters, please take a look to your broker's `message.nack(**extra_options)` method
+    signature.
+
+    Args:
+        extra_options (Any): Additional parameters that will be passed to `message.nack(**extra_options)` method.
+    """
+
+    def __init__(self, **kwargs: Any):
+        self.extra_options = kwargs
+        super().__init__()
 
     def __str__(self) -> str:
         return "Message was nacked"
 
 
 class RejectMessage(HandlerException):
-    """Raise it to `reject` a message immediately."""
+    """Exception raised to reject a message immediately.
+
+    This exception can be used to reject a message with additional options.
+    To watch all allowed parameters, please take a look to your broker's `message.reject(**extra_options)` method
+    signature.
+
+    Args:
+        extra_options (Any): Additional parameters that will be passed to `message.reject(**extra_options)` method.
+    """
+
+    def __init__(self, **kwargs: Any):
+        self.extra_options = kwargs
+        super().__init__()
 
     def __str__(self) -> str:
         return "Message was rejected"

tests/brokers/test_pushback.py

@@ -69,26 +69,6 @@ async def test_push_back_watcher(async_mock: AsyncMock, message):
     message.reject.assert_awaited_once()
 
 
-@pytest.mark.asyncio()
-async def test_push_back_watcher_with_nack_exception(async_mock: AsyncMock, message):
-    watcher = CounterWatcher(3)
-
-    context = WatcherContext(
-        message=message,
-        watcher=watcher,
-    )
-
-    async_mock.side_effect = NackMessage()
-
-    while not message.reject.called and message.nack.await_count < 5:
-        async with context:
-            await async_mock()
-
-    assert not message.ack.await_count
-    assert message.nack.await_count == 3
-    message.reject.assert_awaited_once()
-
-
 @pytest.mark.asyncio()
 async def test_push_endless_back_watcher(async_mock: AsyncMock, message):
     watcher = EndlessWatcher()
@@ -125,3 +105,18 @@ async def test_ignore_skip(async_mock: AsyncMock, message):
     assert not message.nack.called
     assert not message.reject.called
     assert not message.ack.called
+
+
+@pytest.mark.asyncio()
+async def test_additional_params_with_handler_exception(async_mock: AsyncMock, message):
+    watcher = EndlessWatcher()
+
+    context = WatcherContext(
+        message=message,
+        watcher=watcher,
+    )
+
+    async with context:
+        raise NackMessage(delay=5)
+
+    message.nack.assert_called_with(delay=5)