Add new kw params to action_fetch. Fixes #134.

amulet/deployer.py

@@ -15,6 +15,7 @@
     juju,
     timeout as unit_timesout,
     JUJU_VERSION,
+    TimeoutError,
 )
 from .sentry import Talisman
 from .charm import CharmCache
@@ -582,14 +583,27 @@ def action_do(self, unit, action, action_args={}):
         results_id = action_result["Action queued with id"]
         return results_id
 
-    def action_fetch(self, action_id, timeout=600):
+    def action_fetch(
+            self, action_id, timeout=600, raise_on_timeout=False,
+            full_output=False):
         """Fetch results for an action.
 
         If the timeout expires and the action is still not complete, an
-        empty dictionary is returned.
+        empty dictionary is returned. To raise an exception instead, pass
+        ``raise_on_timeout=True``.
+
+        By default, only the 'results' dictionary of the action output is
+        returned. To get the full action output instead, pass
+        ``full_output=True``.
 
         :param action_id: UUID of the action.
         :param timeout: Length of time to wait for an action to complete.
+        :param raise_on_timeout: If True, :class:`amulet.helpers.TimeoutError`
+            will be raised if the action is still running when the timeout
+            expires.
+        :param full_output: If True, returns the full output from the action.
+            If False, only the 'results' dictionary from the action output is
+            returned.
         :return: Action results, as json.
 
         """
@@ -603,9 +617,19 @@ def action_fetch(self, action_id, timeout=600):
         raw = juju(cmd)
         result = json.loads(raw)
         status = result['status']
+
+        if status == 'running' and raise_on_timeout:
+            raise TimeoutError(
+                'Action {} still running after {}s'.format(
+                    action_id, timeout))
+
+        if full_output:
+            return result
+
         if status == 'completed':
             if 'results' in result:
                 return result['results']
+
         return {}
 
     def setup(self, timeout=600, cleanup=True):

tests/test_deployer.py

@@ -10,7 +10,7 @@
 from amulet import Deployment
 from amulet.deployer import get_charm_name
 from amulet.sentry import UnitSentry
-from amulet.helpers import JUJU_VERSION
+from amulet.helpers import JUJU_VERSION, TimeoutError
 from mock import patch, MagicMock, call
 from collections import OrderedDict
 
@@ -563,6 +563,32 @@ def test_action_fetch_wait(self, mj):
                 call(['run-action', 'mysql/0', 'run', '--format', 'json']),
                 call(['show-action-output', 'some-action-id', '--format', 'json', '--wait', '600'])])
 
+    @patch('amulet.deployer.juju')
+    def test_action_fetch_raise_on_timeout(self, mj):
+        mj.side_effect = ['{"Action queued with id": "some-action-id"}',
+                          '{"status":"running",'
+                          '"timing":{"enqueued":"2015-07-21 09:50:59 +0300 EEST",'
+                          '"started":"2015-07-21 09:51:04 +0300 EEST"}}']
+        d = Deployment(juju_env='gojuju')
+        d.add('mysql')
+        uuid = d.action_do('mysql/0', 'run')
+        with self.assertRaises(TimeoutError):
+            d.action_fetch(uuid, timeout=None, raise_on_timeout=True)
+
+    @patch('amulet.deployer.juju')
+    def test_action_fetch_full_output(self, mj):
+        action_output = ['{"Action queued with id": "some-action-id"}',
+                         '{"results":{"key":"value"},"status":"completed",'
+                         '"timing":{"completed":"2015-07-21 09:05:11 +0300 EEST",'
+                         '"enqueued":"2015-07-21 09:05:06 +0300 EEST",'
+                         '"started":"2015-07-21 09:05:09 +0300 EEST"}}']
+        mj.side_effect = action_output
+        d = Deployment(juju_env='gojuju')
+        d.add('mysql')
+        uuid = d.action_do('mysql/0', 'run')
+        results = d.action_fetch(uuid, full_output=True)
+        self.assertEquals(results, json.loads(action_output[1]))
+
     @patch('amulet.deployer.juju')
     def test_unrelate(self, mj):
         d = Deployment(juju_env='gogo')