From b1afefd05b700768813aca783fd5a0c3f5a62632 Mon Sep 17 00:00:00 2001
From: "changjun.zhu" <changjun.zhu@graviti.com>
Date: Tue, 21 Dec 2021 16:08:28 +0800
Subject: [PATCH] feat(client): add update to Job

---
 tensorbay/client/job.py | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/tensorbay/client/job.py b/tensorbay/client/job.py
index 3203ebdcc..cb2a0ed8f 100644
--- a/tensorbay/client/job.py
+++ b/tensorbay/client/job.py
@@ -5,16 +5,21 @@
 
 """Basic structures of asynchronous jobs."""
 
-from typing import Any, Dict, Optional, Tuple, Type, TypeVar
+import time
+from typing import Any, Callable, Dict, Optional, Tuple, Type, TypeVar
 
 from tensorbay.client.struct import Draft
 from tensorbay.utility import AttrsMixin, ReprMixin, ReprType, attr, camel, common_loads
 
+_JOB_UPDATE_INTERVAL = 5
+_JOB_NOT_COMPLETE_STATUS = {"QUEUING", "PROCESSING"}
+
 
 class Job(AttrsMixin, ReprMixin):  # pylint: disable=too-many-instance-attributes
     """This class defines :class:`Job`.
 
     Arguments:
+        job_updater: The function to update the information of the Job instance.
         title: Title of the Job.
         job_id: ID of the Job.
         arguments: Arguments of the Job.
@@ -56,6 +61,7 @@ class Job(AttrsMixin, ReprMixin):  # pylint: disable=too-many-instance-attribute
 
     def __init__(  # pylint: disable=too-many-arguments
         self,
+        job_updater: Callable[[str], Dict[str, Any]],
         title: str,
         job_id: str,
         arguments: Dict[str, Any],
@@ -67,6 +73,7 @@ def __init__(  # pylint: disable=too-many-arguments
         result: Optional[Dict[str, Any]],
         description: Optional[str] = "",
     ) -> None:
+        self._job_updater = job_updater
         self.title = title
         self.job_id = job_id
         self.arguments = arguments
@@ -114,6 +121,21 @@ def update(self, until_complete: bool = False) -> None:
             until_complete: Whether to update job information until it is complete.
 
         """
+        job_info = self._job_updater(self.job_id)
+
+        while until_complete:
+            status = job_info["status"]
+            if status in _JOB_NOT_COMPLETE_STATUS:
+                time.sleep(_JOB_UPDATE_INTERVAL)
+                job_info = self._job_updater(self.job_id)
+                continue
+
+            break
+
+        self.finished_at = job_info["finishedAt"]
+        self.status = job_info["status"]
+        self.error_message = job_info["errorMessage"]
+        self._result = job_info["results"]
 
     def abort(self) -> None:
         """Abort a :class:`Job`."""