From f1c8a47155a703e4056d14cb7fce05dd09248148 Mon Sep 17 00:00:00 2001
From: Stephan Fitzpatrick <knowsuchagency@gmail.com>
Date: Sat, 25 Nov 2023 00:38:53 -0800
Subject: [PATCH] Corrently handle terminal state when invoking job through
 run_script method (#2703)

---
 python-client/wmill/wmill/client.py | 40 ++++++++++++++++-------------
 1 file changed, 22 insertions(+), 18 deletions(-)

diff --git a/python-client/wmill/wmill/client.py b/python-client/wmill/wmill/client.py
index bdbeb902cf..2ce2445815 100644
--- a/python-client/wmill/wmill/client.py
+++ b/python-client/wmill/wmill/client.py
@@ -125,7 +125,7 @@ class Windmill:
             atexit.register(cancel_job)
 
         while True:
-            job = self.get(f"/w/{self.workspace}/jobs_u/get/{job_id}").json()
+            job = self.get_job(job_id)
 
             if timeout and ((time.time() - start_time) > timeout):
                 msg = "reached timeout"
@@ -139,17 +139,23 @@ class Windmill:
             result = job.get("result")
             canceled, canceled_reason = job.get("canceled"), job.get("canceled_reason")
             success = job.get("success")
+            job_type = job.get("type", "")
+            completed = job_type.lower() == "completedjob"
 
-            if cleanup and (canceled or success):
+            if cleanup and completed:
                 atexit.unregister(cancel_job)
 
-            if canceled:
-                raise Exception(f"job canceled: {canceled_reason}")
-
-            if success:
-                if assert_result_is_not_none and result is None:
-                    raise Exception(f"result is None for {job_id = }")
-                return result
+            if completed:
+                if success:
+                    if assert_result_is_not_none and result is None:
+                        raise Exception(f"result is None for {job_id = }")
+                    return result
+                else:
+                    if canceled:
+                        raise Exception(f"job canceled: {canceled_reason}")
+                    else:
+                        error = result.get("error")
+                        raise Exception(f"job failed: {error}")
 
             if verbose:
                 logger.info(f"sleeping 0.5 seconds for {job_id = }")
@@ -189,18 +195,16 @@ class Windmill:
 
         return result
 
+    def get_job(self, job_id: str) -> dict:
+        return self.get(f"/w/{self.workspace}/jobs_u/get/{job_id}").json()
+
     def get_job_status(self, job_id: str) -> JobStatus:
-        resp = self.get(
-            f"/w/{self.workspace}/jobs_u/get/{job_id}",
-            raise_for_status=False,
-        )
-        assert not resp.status_code == 404, f"{job_id} not found"
-        resp_json = resp.json()
-        job_type = resp_json.get("type", "")
-        assert job_type, f"{resp_json} is not a valid job"
+        job = self.get_job(job_id)
+        job_type = job.get("type", "")
+        assert job_type, f"{job} is not a valid job"
         if job_type.lower() == "completedjob":
             return "COMPLETED"
-        additional_properties = resp_json.get("additional_properties", {})
+        additional_properties = job.get("additional_properties", {})
         if "running" not in additional_properties:
             raise Exception(f"{job_id} is not running")
         if additional_properties.get("running"):