Repository: botify-labs/simpleflow
Issue: activity_rerun fails on MetrologyTask's

The `MetrologyTask.can_upload` method looks for keys in `self.context` without first checking whether this attribute is set; and `activity_rerun` doesn't set it.

Changed files:
tests/test_simpleflow/test_helpers.py, simpleflow/command.py, simpleflow/swf/helpers.py, simpleflow/metrology.py

diff --git a/simpleflow/command.py b/simpleflow/command.py
index 511a2e38..644f97ce 100644
--- a/simpleflow/command.py
+++ b/simpleflow/command.py
@@ -686,6 +686,10 @@ def activity_rerun(domain,
     task, args, kwargs, meta, params = helpers.find_activity(
         history, scheduled_id=scheduled_id, activity_id=activity_id, input=input_override,
     )
+    kwargs["context"].update({
+        "workflow_id": wfe.workflow_id,
+        "run_id": wfe.run_id,
+    })
     logger.debug("Found activity. Last execution:")
     for line in json_dumps(params, pretty=True).split("\\n"):
         logger.debug(line)
diff --git a/simpleflow/metrology.py b/simpleflow/metrology.py
index 1a24d1a9..f4814e64 100644
--- a/simpleflow/metrology.py
+++ b/simpleflow/metrology.py
@@ -84,6 +84,8 @@ class StepExecution(object):
 class MetrologyTask(object):
 
     def can_upload(self):
+        if not hasattr(self, 'context'):
+            return False
         return all(c in self.context for c in ('workflow_id', 'run_id', 'activity_id'))
 
     @property
diff --git a/simpleflow/swf/helpers.py b/simpleflow/swf/helpers.py
index b07f2af0..89632823 100644
--- a/simpleflow/swf/helpers.py
+++ b/simpleflow/swf/helpers.py
@@ -1,12 +1,13 @@
 from __future__ import absolute_import
 
+import copy
 import getpass
 import json
 import os
 import socket
 
 import psutil
-from future.utils import iteritems
+from future.utils import iteritems, itervalues
 
 import swf.exceptions
 import swf.models
@@ -107,7 +108,7 @@ def find_activity(history, scheduled_id=None, activity_id=None, input=None):
     :type input: Optional[dict[str, Any]]
     """
     found_activity = None
-    for _, params in iteritems(history.activities):
+    for params in itervalues(history.activities):
         if params["scheduled_id"] == scheduled_id:
             found_activity = params
         if params["id"] == activity_id:
@@ -129,6 +130,13 @@ def find_activity(history, scheduled_id=None, activity_id=None, input=None):
     kwargs = input_.get('kwargs', {})
     meta = input_.get('meta', {})
 
+    kwargs["context"] = {
+        "name": activity_str,
+        "version": None,  # not stored in the activity dict
+        "activity_id": found_activity["id"],
+        "input": copy.deepcopy(input_),
+    }
+
     # return everything
     return activity, args, kwargs, meta, found_activity
 
diff --git a/tests/test_simpleflow/test_helpers.py b/tests/test_simpleflow/test_helpers.py
index b1d6a8bd..5166ce2b 100644
--- a/tests/test_simpleflow/test_helpers.py
+++ b/tests/test_simpleflow/test_helpers.py
@@ -26,7 +26,16 @@ class TestSimpleflowSwfHelpers(unittest.TestCase):
         func, args, kwargs, meta, params = find_activity(FakeHistory(), scheduled_id=5)
         expect(str(func)).to.match(r'^Activity\\(name=tests.integration.workflow.sleep,')
         expect(args).to.equal([37])
-        expect(kwargs).to.equal({})
+        expect(kwargs).to.equal(
+            {
+                'context': {
+                    'activity_id': 'activity-tests.integration.workflow.sleep-1',
+                    'input': {'args': [37]},
+                    'name': 'tests.integration.workflow.sleep',
+                    'version': None,
+                }
+            }
+        )
         expect(meta).to.equal({})
         expect(params["id"]).to.equal("activity-tests.integration.workflow.sleep-1")