diff --git a/moto/swf/models/domain.py b/moto/swf/models/domain.py index 4ed91452..d04a4184 100644 --- a/moto/swf/models/domain.py +++ b/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 diff --git a/moto/swf/models/workflow_execution.py b/moto/swf/models/workflow_execution.py index b0e6427c..a9a27725 100644 --- a/moto/swf/models/workflow_execution.py +++ b/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, diff --git a/tests/test_swf/models/test_domain.py b/tests/test_swf/models/test_domain.py index 68e3f890..5d2982e1 100644 --- a/tests/test_swf/models/test_domain.py +++ b/tests/test_swf/models/test_domain.py @@ -21,9 +21,28 @@ def test_domain_string_representation(): domain = Domain("my-domain", "60") str(domain).should.equal("Domain(name: my-domain, status: REGISTERED)") -def test_domain_add_to_task_list(): +def test_domain_add_to_activity_task_list(): domain = Domain("my-domain", "60") - domain.add_to_task_list("foo", "bar") - dict(domain.task_lists).should.equal({ + domain.add_to_activity_task_list("foo", "bar") + domain.activity_task_lists.should.equal({ "foo": ["bar"] }) + +def test_domain_activity_tasks(): + domain = Domain("my-domain", "60") + domain.add_to_activity_task_list("foo", "bar") + domain.add_to_activity_task_list("other", "baz") + domain.activity_tasks.should.equal(["bar", "baz"]) + +def test_domain_add_to_decision_task_list(): + domain = Domain("my-domain", "60") + domain.add_to_decision_task_list("foo", "bar") + domain.decision_task_lists.should.equal({ + "foo": ["bar"] + }) + +def test_domain_decision_tasks(): + domain = Domain("my-domain", "60") + domain.add_to_decision_task_list("foo", "bar") + domain.add_to_decision_task_list("other", "baz") + domain.decision_tasks.should.equal(["bar", "baz"]) diff --git a/tests/test_swf/models/test_workflow_execution.py b/tests/test_swf/models/test_workflow_execution.py index b793d8d2..c6660ab9 100644 --- a/tests/test_swf/models/test_workflow_execution.py +++ b/tests/test_swf/models/test_workflow_execution.py @@ -207,7 +207,7 @@ def test_workflow_execution_schedule_activity_task(): task = wfe.activity_tasks[0] task.activity_id.should.equal("my-activity-001") task.activity_type.name.should.equal("test-activity") - wfe.domain.task_lists["task-list-name"].should.contain(task) + wfe.domain.activity_task_lists["task-list-name"].should.contain(task) def test_workflow_execution_schedule_activity_task_without_task_list_should_take_default(): wfe = make_workflow_execution() @@ -229,7 +229,7 @@ def test_workflow_execution_schedule_activity_task_without_task_list_should_take last_event.task_list.should.equal("foobar") task = wfe.activity_tasks[0] - wfe.domain.task_lists["foobar"].should.contain(task) + wfe.domain.activity_task_lists["foobar"].should.contain(task) def test_workflow_execution_schedule_activity_task_should_fail_if_wrong_attributes(): wfe = make_workflow_execution() @@ -286,7 +286,7 @@ def test_workflow_execution_schedule_activity_task_should_fail_if_wrong_attribut wfe.open_counts["openActivityTasks"].should.equal(0) wfe.activity_tasks.should.have.length_of(0) - wfe.domain.task_lists.should.have.length_of(0) + wfe.domain.activity_task_lists.should.have.length_of(0) hsh["heartbeatTimeout"] = "300" wfe.schedule_activity_task(123, hsh) @@ -294,7 +294,7 @@ def test_workflow_execution_schedule_activity_task_should_fail_if_wrong_attribut last_event.event_type.should.equal("ActivityTaskScheduled") task = wfe.activity_tasks[0] - wfe.domain.task_lists["foobar"].should.contain(task) + wfe.domain.activity_task_lists["foobar"].should.contain(task) wfe.open_counts["openDecisionTasks"].should.equal(0) wfe.open_counts["openActivityTasks"].should.equal(1)