44import pytest
55from aio_pika import Channel , Message
66from aio_pika .exceptions import QueueEmpty
7- from taskiq import BrokerMessage
7+ from taskiq import AckableMessage , BrokerMessage
8+ from taskiq .utils import maybe_awaitable
89
910from taskiq_aio_pika .broker import AioPikaBroker
1011
1112
12- async def get_first_task (broker : AioPikaBroker ) -> bytes : # type: ignore
13+ async def get_first_task (broker : AioPikaBroker ) -> AckableMessage : # type: ignore
1314 """
1415 Get first message from the queue.
1516
@@ -46,7 +47,8 @@ async def test_kick_success(broker: AioPikaBroker) -> None:
4647
4748 message = await asyncio .wait_for (get_first_task (broker ), timeout = 0.4 )
4849
49- assert message == sent .message
50+ assert message .data == sent .message
51+ await maybe_awaitable (message .ack ())
5052
5153
5254@pytest .mark .anyio
@@ -111,7 +113,8 @@ async def test_listen(
111113
112114 message = await asyncio .wait_for (get_first_task (broker ), timeout = 0.4 )
113115
114- assert message == b"test_message"
116+ assert message .data == b"test_message"
117+ await maybe_awaitable (message .ack ())
115118
116119
117120@pytest .mark .anyio
@@ -133,9 +136,10 @@ async def test_wrong_format(
133136 routing_key = queue_name ,
134137 )
135138
136- msg_bytes = await asyncio .wait_for (get_first_task (broker ), 0.4 )
139+ message = await asyncio .wait_for (get_first_task (broker ), 0.4 )
137140
138- assert msg_bytes == b"wrong"
141+ assert message .data == b"wrong"
142+ await maybe_awaitable (message .ack ())
139143
140144 with pytest .raises (QueueEmpty ):
141145 await queue .get ()
0 commit comments