Merge pull request #11 from taskiq-python/bugfix/ackable-msg

Author: s3rius

poetry.lock

@@ -21,7 +21,7 @@ keywords = ["taskiq", "tasks", "distributed", "async", "nats", "result_backend"]
 [tool.poetry.dependencies]
 python = "^3.8.1"
 nats-py = "^2.2.0"
-taskiq = "^0"
+taskiq = ">=0.8,<1"
 
 
 [tool.poetry.group.dev.dependencies]

pyproject.toml

@@ -4,7 +4,7 @@
 from nats.aio.client import Client
 from nats.js import JetStreamContext
 from nats.js.api import StreamConfig
-from taskiq import AsyncBroker, AsyncResultBackend, BrokerMessage
+from taskiq import AckableMessage, AsyncBroker, AsyncResultBackend, BrokerMessage
 
 _T = TypeVar("_T")  # noqa: WPS111 (Too short)
 
@@ -140,16 +140,18 @@ async def kick(self, message: BrokerMessage) -> None:
             headers=message.labels,
         )
 
-    async def listen(self) -> AsyncGenerator[bytes, None]:
+    async def listen(self) -> AsyncGenerator[AckableMessage, None]:
         """
         Start listen to new messages.
 
         :yield: incoming messages.
         """
         subscribe = await self.js.subscribe(self.subject, queue=self.queue or "")
         async for message in subscribe.messages:
-            yield message.data
-            await message.ack()
+            yield AckableMessage(
+                data=message.data,
+                ack=message.ack,
+            )
 
     async def shutdown(self) -> None:
         """Close connections to NATS."""

taskiq_nats/broker.py

@@ -3,7 +3,7 @@
 from typing import List
 
 import pytest
-from taskiq import BrokerMessage
+from taskiq import AckableMessage, BrokerMessage
 
 from taskiq_nats import JetStreamBroker
 from tests.utils import read_message
@@ -33,5 +33,10 @@ async def test_success(nats_urls: List[str], nats_subject: str) -> None:
         labels={},
     )
     await broker.kick(sent_message)
-    assert await asyncio.wait_for(read_message(broker), 0.5) == sent_message.message
+    ackable_msg = await asyncio.wait_for(read_message(broker), 0.5)
+    assert isinstance(ackable_msg, AckableMessage)
+    assert ackable_msg.data == sent_message.message
+    ack = ackable_msg.ack()
+    if ack is not None:
+        await ack
     await broker.shutdown()

tests/test_jstream.py

@@ -1,13 +1,17 @@
-from taskiq import AsyncBroker
+from typing import Union
 
+from taskiq import AckableMessage, AsyncBroker
 
-async def read_message(broker: AsyncBroker) -> bytes:  # type: ignore
+
+async def read_message(broker: AsyncBroker) -> Union[bytes, AckableMessage]:
     """
     Read signle message from the broker's listen method.
 
     :param broker: current broker.
     :return: firs message.
     """
+    msg: Union[bytes, AckableMessage] = b"error"
     async for message in broker.listen():
-        assert isinstance(message, bytes)  # noqa: S101
-        return message
+        msg = message
+        break
+    return msg