From 3bac93677f5c1bd4d33896e13fb25642cdf5bb7b Mon Sep 17 00:00:00 2001 From: Chris Harris Date: Thu, 26 Sep 2024 21:04:12 +0000 Subject: [PATCH] Use status to determine exit condition As the result object of a task can now be present while the task is running we need to use the status to wait for task completion. --- src/sfapi_client/_async/compute.py | 5 +++-- src/sfapi_client/_sync/compute.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/sfapi_client/_async/compute.py b/src/sfapi_client/_async/compute.py index 8a476c5..fbfb827 100644 --- a/src/sfapi_client/_async/compute.py +++ b/src/sfapi_client/_async/compute.py @@ -58,10 +58,11 @@ async def _wait_for_task(self, task_id) -> str: json_response = r.json() task = Task.model_validate(json_response) - if task.status.lower() in ["error", "failed"]: + status = task.status.lower() + if status in ["error", "failed"]: raise SfApiError(task.result) - if task.result is None: + if status not in ["completed", "cancelled"]: await _ASYNC_SLEEP(1) continue diff --git a/src/sfapi_client/_sync/compute.py b/src/sfapi_client/_sync/compute.py index 0c7ca0b..db437c8 100644 --- a/src/sfapi_client/_sync/compute.py +++ b/src/sfapi_client/_sync/compute.py @@ -58,10 +58,11 @@ def _wait_for_task(self, task_id) -> str: json_response = r.json() task = Task.model_validate(json_response) - if task.status.lower() in ["error", "failed"]: + status = task.status.lower() + if status in ["error", "failed"]: raise SfApiError(task.result) - if task.result is None: + if status not in ["completed", "cancelled"]: _SLEEP(1) continue