Ensure we document and complete correct removal.

Change-Id: I156a7d53c3640f6be64532a5c6b74f71ed651d14

taskflow/patterns/graph_flow.py

@@ -65,16 +65,17 @@ class Flow(linear_flow.Flow):
         return "; ".join(lines)
 
     @decorators.locked
-    def remove(self, task_uuid):
+    def remove(self, uuid):
-        remove_nodes = []
+        runner = None
         for r in self._graph.nodes_iter():
-            if r.uuid == task_uuid:
+            if r.uuid == uuid:
-                remove_nodes.append(r)
+                runner = r
-        if not remove_nodes:
+                break
-            raise IndexError("No task found with uuid %s" % (task_uuid))
+        if not runner:
+            raise ValueError("No runner found with uuid %s" % (uuid))
         else:
-            for r in remove_nodes:
+            # Ensure that we reset out internal state after said removal
-                self._graph.remove_node(r)
+            self._graph.remove_node(runner)
             self._runners = []
             self._leftoff_at = None
 

taskflow/patterns/linear_flow.py

@@ -103,17 +103,26 @@ class Flow(base.Flow):
         return "; ".join(lines)
 
     @decorators.locked
-    def remove(self, task_uuid):
+    def remove(self, uuid):
-        removed = False
+        index_removed = -1
         for (i, r) in enumerate(self._runners):
-            if r.uuid == task_uuid:
+            if r.uuid == uuid:
-                self._runners.pop(i)
+                index_removed = i
+                break
+        if index_removed == -1:
+            raise ValueError("No runner found with uuid %s" % (uuid))
+        else:
+            # Ensure that we reset out internal state after said removal.
+            removed = self._runners.pop(index_removed)
             self._connected = False
             self._leftoff_at = None
-                removed = True
+            # Go and remove it from any runner after the removed runner since
-                break
+            # those runners may have had an attachment to it.
-        if not removed:
+            for r in self._runners[index_removed:]:
-            raise IndexError("No task found with uuid %s" % (task_uuid))
+                try:
+                    r.runs_before.remove(removed)
+                except (IndexError, ValueError):
+                    pass
 
     def _connect(self):
         if self._connected: