From 78b38eb70e6ba9c8699f96b159bf25b94e50cbb1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nguy=E1=BB=85n=20Anh=20B=C3=ACnh?=
 <sometimesocrazy@gmail.com>
Date: Mon, 12 Jun 2023 02:06:42 +0700
Subject: [PATCH] - Fix bug monitoring

---
 evoflow/VERSION                    |  2 +-
 evoflow/entities/core/job.py       | 31 ++++++++++++++++++++----------
 evoflow/entities/core/step.py      | 23 ++++++++++++++--------
 evoflow/entities/core/step_list.py | 30 ++++++++++++++++++++++++++++-
 4 files changed, 66 insertions(+), 20 deletions(-)

diff --git a/evoflow/VERSION b/evoflow/VERSION
index 756bb3a..2887d21 100644
--- a/evoflow/VERSION
+++ b/evoflow/VERSION
@@ -1 +1 @@
-0.1.9.1
+0.1.9.2
diff --git a/evoflow/entities/core/job.py b/evoflow/entities/core/job.py
index 36f6caa..da5bc21 100644
--- a/evoflow/entities/core/job.py
+++ b/evoflow/entities/core/job.py
@@ -249,17 +249,28 @@ def remove_running_step(self, step):
 
     def __update_live(self, live):
         tree = Tree(self.name)
-        for step in self.__running_steps:
-            is_step_list = isinstance(step, StepList)
-            if is_step_list:
-                remaining_steps = len(step.get_remaining_step())
-                total_step = len(step.steps)
-                step_title = f"{step.name} {total_step-remaining_steps}/{total_step}"
-            else:
-                step_title = step.name
+        running_steps = self.__running_steps
+
+        tree_added_steps = []
+
+        step_lists = [_ for _ in running_steps if isinstance(_, StepList)]
+        single_steps = [_ for _ in running_steps if not isinstance(_, Step)]
+
+        for step_list in step_lists:
+            remaining_steps = len(step_list.get_remaining_step())
+            total_step = len(step_list.steps)
+            step_title = f"{step_list.name} {total_step-remaining_steps}/{total_step}"
+
             spinner = Spinner("material", text=Text(step_title, style="green"))
             step_live = tree.add(spinner)
-            if isinstance(step, StepList):
-                for sub_step in step.steps:
+            for sub_step in step_list.steps:
+                if sub_step.is_running():
                     step_live.add(Spinner("material", text=Text(sub_step.name, style="blue")))
+                tree_added_steps.append(sub_step)
+        for step in single_steps:
+            if step in tree_added_steps:
+                continue
+            if step.is_running():
+                tree.add(Spinner("material", text=Text(step.name, style="blue")))
+
         live.update(Panel(tree, title=f"Running {self.name}"))
diff --git a/evoflow/entities/core/step.py b/evoflow/entities/core/step.py
index 8f99ec5..c7c682f 100644
--- a/evoflow/entities/core/step.py
+++ b/evoflow/entities/core/step.py
@@ -70,9 +70,17 @@ def __init__(self, name=None, transactions=None, **kwargs):
         self.previous_steps = []
         self.error = None
         self.status = self.STATUS_PENDING
-        self.job = None
+        self._job = None
         super().__init__(name=name, **kwargs)
 
+    @property
+    def job(self):
+        return self._job
+
+    @job.setter
+    def job(self, value):
+        self._job = value
+
     @abc.abstractmethod
     def action(self, **kwargs):
         """
@@ -80,13 +88,6 @@ def action(self, **kwargs):
         """
         pass
 
-    def do_action(self, **kwargs):
-        """
-        Performs the function of step
-        """
-
-        return self.action(**kwargs)
-
     def set_error(self, error):
         self.status = self.STATUS_FAILED
         self.error = error
@@ -160,3 +161,9 @@ def is_ready(self):
                 return False
             self.status = self.STATUS_READY
         return True
+
+    def is_running(self):
+        """
+        Check if step is running
+        """
+        return self.status == self.STATUS_RUNNING
\ No newline at end of file
diff --git a/evoflow/entities/core/step_list.py b/evoflow/entities/core/step_list.py
index b77a729..4f4cdb3 100644
--- a/evoflow/entities/core/step_list.py
+++ b/evoflow/entities/core/step_list.py
@@ -40,6 +40,24 @@ def __str__(self):
     def __getitem__(self, item):
         return self.steps[item]
 
+    def prepare(self, **kwargs):
+        """
+        Prepare the step to run
+        """
+
+        for step in self.steps:
+            step.prepare(**kwargs)
+        return super().prepare(**kwargs)
+
+    def end(self, **kwargs) -> dict:
+        """
+        End the step
+        """
+
+        for step in self.steps:
+            step.end(**kwargs)
+        return super().end(**kwargs)
+
     def action(self, **kwargs):
         """
         Performs the function of step
@@ -49,4 +67,14 @@ def action(self, **kwargs):
             list(executor.map(lambda step: step.action(**kwargs), self.steps))
 
     def get_remaining_step(self):
-        return [step for step in self.steps if step.status != Step.STATUS_SUCCESS]
+        return [step for step in self.steps if not step.is_running()]
+
+    @property
+    def job(self):
+        return super().job
+
+    @job.setter
+    def job(self, value):
+        self._job = value
+        for step in self.steps:
+            step.job = value