Merge "Allow for revert to have a different argument list from execute"

3 files changed, 21 insertions, 10 deletions

diff --git a/taskflow/engines/action_engine/actions/retry.py b/taskflow/engines/action_engine/actions/retry.py
index 126b903..49dd94a 100644
--- a/taskflow/engines/action_engine/actions/retry.py
+++ b/taskflow/engines/action_engine/actions/retry.py
@@ -30,12 +30,19 @@ class RetryAction(base.Action):
         super(RetryAction, self).__init__(storage, notifier)
         self._retry_executor = retry_executor
 
-    def _get_retry_args(self, retry, addons=None):
-        arguments = self._storage.fetch_mapped_args(
-            retry.rebind,
-            atom_name=retry.name,
-            optional_args=retry.optional
-        )
+    def _get_retry_args(self, retry, revert=False, addons=None):
+        if revert:
+            arguments = self._storage.fetch_mapped_args(
+                retry.revert_rebind,
+                atom_name=retry.name,
+                optional_args=retry.revert_optional
+            )
+        else:
+            arguments = self._storage.fetch_mapped_args(
+                retry.rebind,
+                atom_name=retry.name,
+                optional_args=retry.optional
+            )
         history = self._storage.get_retry_history(retry.name)
         arguments[retry_atom.EXECUTE_REVERT_HISTORY] = history
         if addons:
@@ -92,7 +99,7 @@ class RetryAction(base.Action):
             retry_atom.REVERT_FLOW_FAILURES: self._storage.get_failures(),
         }
         return self._retry_executor.revert_retry(
-            retry, self._get_retry_args(retry, addons=arg_addons))
+            retry, self._get_retry_args(retry, addons=arg_addons, revert=True))
 
     def on_failure(self, retry, atom, last_failure):
         self._storage.save_retry_failure(retry.name, atom.name, last_failure)
diff --git a/taskflow/engines/action_engine/actions/task.py b/taskflow/engines/action_engine/actions/task.py
index ac117e1..6d0981d 100644
--- a/taskflow/engines/action_engine/actions/task.py
+++ b/taskflow/engines/action_engine/actions/task.py
@@ -122,9 +122,9 @@ class TaskAction(base.Action):
     def schedule_reversion(self, task):
         self.change_state(task, states.REVERTING, progress=0.0)
         arguments = self._storage.fetch_mapped_args(
-            task.rebind,
+            task.revert_rebind,
             atom_name=task.name,
-            optional_args=task.optional
+            optional_args=task.revert_optional
         )
         task_uuid = self._storage.get_atom_uuid(task.name)
         task_result = self._storage.get(task.name)
diff --git a/taskflow/engines/action_engine/engine.py b/taskflow/engines/action_engine/engine.py
index d09e7b3..f7625a7 100644
--- a/taskflow/engines/action_engine/engine.py
+++ b/taskflow/engines/action_engine/engine.py
@@ -379,8 +379,12 @@ class ActionEngine(base.Engine):
         last_node = None
         missing_nodes = 0
         for atom in self._runtime.analyzer.iterate_nodes(compiler.ATOMS):
-            atom_missing = self.storage.fetch_unsatisfied_args(
+            exec_missing = self.storage.fetch_unsatisfied_args(
                 atom.name, atom.rebind, optional_args=atom.optional)
+            revert_missing = self.storage.fetch_unsatisfied_args(
+                atom.name, atom.revert_rebind,
+                optional_args=atom.revert_optional)
+            atom_missing = exec_missing.union(revert_missing)
             if atom_missing:
                 cause = exc.MissingDependencies(atom,
                                                 sorted(atom_missing),