diff --git a/aiojobs/_job.py b/aiojobs/_job.py index e8e2102..41f3d34 100644 --- a/aiojobs/_job.py +++ b/aiojobs/_job.py @@ -165,3 +165,6 @@ def _report_exception(self, exc: BaseException) -> None: if self._source_traceback is not None: context["source_traceback"] = self._source_traceback self._scheduler.call_exception_handler(context) + + def __await__(self) -> _T: + return self.wait().__await__() diff --git a/tests/test_job.py b/tests/test_job.py index b85f8d7..5c9fe3d 100644 --- a/tests/test_job.py +++ b/tests/test_job.py @@ -314,3 +314,23 @@ async def coro() -> None: assert job.get_name() == "changed_name" assert job._task is not None assert job._task.get_name() == "changed_name" + + +async def test_awaitable_job(scheduler: Scheduler) -> None: + async def coro() -> int: + """Dummy function with result""" + return 1 + + job = await scheduler.spawn(coro()) + result = await job + assert result == 1 + + +async def test_job_gathering(scheduler: Scheduler): + async def coro() -> int: + """Dummy function with result""" + return 1 + + jobs = (scheduler.spawn(coro()) for _ in range(5)) + results = await asyncio.gather(*jobs) + assert results == [1, 1, 1, 1, 1]