Rework task lists for activity/decision tasks

moto/swf/models/domain.py

@@ -23,7 +23,8 @@ class Domain(object):
         # of "workflow_id (client determined)" => WorkflowExecution()
         # here.
         self.workflow_executions = {}
-        self.task_lists = defaultdict(list)
+        self.activity_task_lists = {}
+        self.decision_task_lists = {}
 
     def __repr__(self):
         return "Domain(name: %(name)s, status: %(status)s)" % self.__dict__
@@ -85,5 +86,26 @@ class Domain(object):
             )
         return wfe
 
-    def add_to_task_list(self, task_list, obj):
-        self.task_lists[task_list].append(obj)
+    def add_to_activity_task_list(self, task_list, obj):
+        if not task_list in self.activity_task_lists:
+            self.activity_task_lists[task_list] = []
+        self.activity_task_lists[task_list].append(obj)
+
+    @property
+    def activity_tasks(self):
+        _all = []
+        for _, tasks in self.activity_task_lists.iteritems():
+            _all += tasks
+        return _all
+
+    def add_to_decision_task_list(self, task_list, obj):
+        if not task_list in self.decision_task_lists:
+            self.decision_task_lists[task_list] = []
+        self.decision_task_lists[task_list].append(obj)
+
+    @property
+    def decision_tasks(self):
+        _all = []
+        for _, tasks in self.decision_task_lists.iteritems():
+            _all += tasks
+        return _all

moto/swf/models/workflow_execution.py

@@ -75,9 +75,7 @@ class WorkflowExecution(object):
         }
         # events
         self._events = []
-        # tasks
-        self.decision_tasks = []
-        self.activity_tasks = []
+        # child workflows
         self.child_workflow_executions = []
 
     def __repr__(self):
@@ -167,12 +165,22 @@ class WorkflowExecution(object):
         self.schedule_decision_task()
 
     def schedule_decision_task(self):
-        self.open_counts["openDecisionTasks"] += 1
         evt = self._add_event(
             "DecisionTaskScheduled",
             workflow_execution=self,
         )
-        self.decision_tasks.append(DecisionTask(self, evt.event_id))
+        self.domain.add_to_decision_task_list(
+            self.task_list,
+            DecisionTask(self, evt.event_id),
+        )
+        self.open_counts["openDecisionTasks"] += 1
+
+    @property
+    def decision_tasks(self):
+        return filter(
+            lambda t: t.workflow_execution == self,
+            self.domain.decision_tasks
+        )
 
     @property
     def scheduled_decision_tasks(self):
@@ -181,6 +189,13 @@ class WorkflowExecution(object):
             self.decision_tasks
         )
 
+    @property
+    def activity_tasks(self):
+        return filter(
+            lambda t: t.workflow_execution == self,
+            self.domain.activity_tasks
+        )
+
     def _find_decision_task(self, task_token):
         for dt in self.decision_tasks:
             if dt.task_token == task_token:
@@ -395,9 +410,7 @@ class WorkflowExecution(object):
             workflow_execution=self,
         )
         # Only add event and increment counters if nothing went wrong
-        # TODO: don't store activity tasks in 2 places...
-        self.activity_tasks.append(task)
-        self.domain.add_to_task_list(task_list, task)
+        self.domain.add_to_activity_task_list(task_list, task)
         self._add_event(
             "ActivityTaskScheduled",
             decision_task_completed_event_id=event_id,