From 631be1a68d7e66fb92c886be7e594bc4b234b4e4 Mon Sep 17 00:00:00 2001 From: Andrei Neagu Date: Mon, 13 Jan 2025 17:18:15 +0100 Subject: [PATCH] refactor --- .../fastapi/long_running_tasks/_context_manager.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/packages/service-library/src/servicelib/fastapi/long_running_tasks/_context_manager.py b/packages/service-library/src/servicelib/fastapi/long_running_tasks/_context_manager.py index 35b734ac055..1b0f47449b3 100644 --- a/packages/service-library/src/servicelib/fastapi/long_running_tasks/_context_manager.py +++ b/packages/service-library/src/servicelib/fastapi/long_running_tasks/_context_manager.py @@ -104,17 +104,16 @@ async def _status_update() -> TaskStatus: ) return task_status - async def _wait_task_completion() -> None: + async def _wait_for_task_result() -> Any: task_status = await _status_update() while not task_status.done: await asyncio.sleep(status_poll_interval) task_status = await _status_update() - try: - await asyncio.wait_for(_wait_task_completion(), timeout=task_timeout) - - result = await client.get_task_result(task_id) + return await client.get_task_result(task_id) + try: + result = await asyncio.wait_for(_wait_for_task_result(), timeout=task_timeout) logger.debug("%s, %s", f"{task_id=}", f"{result=}") yield result