Added utilities method for inmemory broker.

Author: s3rius

taskiq/brokers/inmemory_broker.py

@@ -127,6 +127,7 @@ def __init__(
         cast_types: bool = True,
         max_async_tasks: int = 30,
         propagate_exceptions: bool = True,
+        await_inplace: bool = False,
     ) -> None:
         super().__init__()
         self.result_backend = InmemoryResultBackend(
@@ -140,6 +141,7 @@ def __init__(
             max_async_tasks=max_async_tasks,
             propagate_exceptions=propagate_exceptions,
         )
+        self.await_inplace = await_inplace
         self._running_tasks: "Set[asyncio.Task[Any]]" = set()
 
     async def kick(self, message: BrokerMessage) -> None:
@@ -156,7 +158,12 @@ async def kick(self, message: BrokerMessage) -> None:
         if target_task is None:
             raise TaskiqError("Unknown task.")
 
-        task = asyncio.create_task(self.receiver.callback(message=message.message))
+        receiver_cb = self.receiver.callback(message=message.message)
+        if self.await_inplace:
+            await receiver_cb
+            return
+
+        task = asyncio.create_task(receiver_cb)
         self._running_tasks.add(task)
         task.add_done_callback(self._running_tasks.discard)
 
@@ -171,6 +178,17 @@ def listen(self) -> AsyncGenerator[bytes, None]:
         """
         raise RuntimeError("Inmemory brokers cannot listen.")
 
+    async def wait_all(self) -> None:
+        """
+        Wait for all currently running tasks to complete.
+
+        Useful when used in testing and you need to await all sent tasks
+        before asserting results.
+        """
+        to_await = list(self._running_tasks)
+        for task in to_await:
+            await task
+
     async def startup(self) -> None:
         """Runs startup events for client and worker side."""
         for event in (TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.WORKER_STARTUP):

tests/brokers/test_inmemory.py

@@ -85,3 +85,39 @@ async def test_task() -> str:
 
     result = await task.wait_result()
     assert result.return_value == test_value
+
+
+@pytest.mark.anyio
+async def test_inline_awaits() -> None:
+    broker = InMemoryBroker(await_inplace=True)
+    slept = False
+
+    @broker.task
+    async def test_task() -> None:
+        nonlocal slept
+        await asyncio.sleep(0.2)
+        slept = True
+
+    task = await test_task.kiq()
+    assert slept
+    assert await task.is_ready()
+    assert not broker._running_tasks
+
+
+@pytest.mark.anyio
+async def test_wait_all() -> None:
+    broker = InMemoryBroker()
+    slept = False
+
+    @broker.task
+    async def test_task() -> None:
+        nonlocal slept
+        await asyncio.sleep(0.2)
+        slept = True
+
+    task = await test_task.kiq()
+    assert not slept
+    await broker.wait_all()
+    assert slept
+    assert await task.is_ready()
+    assert not broker._running_tasks