Added ability to specify return codes.

provy/core/errors.py

@@ -1,6 +1,10 @@
-#!/usr/bin/python
 # -*- coding: utf-8 -*-
 
 
 class ConfigurationError(RuntimeError):
     '''Raised when there's a configuration error in the provyfile.'''
+
+class CommandExecutionError(RuntimeError):
+    '''
+    Raised when local command has invalid exitcode
+    '''

provy/core/roles.py

@@ -16,6 +16,7 @@
 from jinja2 import Environment, PackageLoader, FileSystemLoader
 import uuid
 from StringIO import StringIO
+from provy.core.errors import CommandExecutionError
 
 
 class UsingRole(object):
@@ -227,7 +228,7 @@ def __cd(self, cd=None):
         else:
             yield
 
-    def execute(self, command, stdout=True, sudo=False, user=None, cwd=None):
+    def execute(self, command, stdout=True, sudo=False, user=None, cwd=None, ok_returncodes = (0, )):
         '''
         This method is the bread and butter of provy and is a base for most other methods that interact with remote servers.
 
@@ -248,6 +249,8 @@ def execute(self, command, stdout=True, sudo=False, user=None, cwd=None):
              cd into that directory before executing command. Current path will be
              *unchanged* after the call.
         :type cwd: :class:`str`
+        :param ok_returncodes: List of returncodes that command may return, and that are not treated as an error.
+        :type ok_returncodes: :class:`list` of :clas:`int`.
 
         :return: The execution result
         :rtype: :class:`str`
@@ -262,16 +265,27 @@ def provision(self):
                     self.execute('ls /', stdout=False, sudo=True)
                     self.execute('ls /', stdout=False, user='vip')
         '''
-        with self.__showing_command_output(stdout):
-            with self.__cd(cwd):
-                return self.__execute_command(command, sudo=sudo, user=user)
+        with self.__showing_command_output(stdout),  self.__cd(cwd), fabric.api.settings(warn_only=True):
+            result =  self.__execute_command(command, sudo=sudo, user=user)
+        self.__check_returncodes(command, result, ok_returncodes)
+        return result
 
     def __execute_command(self, command, sudo=False, user=None):
         if sudo or (user is not None):
             return fabric.api.sudo(command, user=user)
         return fabric.api.run(command)
 
-    def execute_local(self, command, stdout=True, sudo=False, user=None):
+    def __check_returncodes(self, command, result, ok_returncodes):
+        if not result.return_code in ok_returncodes:
+            msg = (
+                "Command '{cmd}' returnd returncode '{ret}' which signifies "
+                "fatal error. Additional information passed to log"
+            )
+            self.log(result)
+            raise CommandExecutionError(msg.format(cmd=command, ret=result.return_code))
+
+
+    def execute_local(self, command, stdout=True, sudo=False, user=None, ok_returncodes = (0, )):
         '''
         Allows you to perform any shell action in the local machine. It is an abstraction over the `fabric.api.local <https://fabric.readthedocs.org/en/latest/api/core/operations.html#fabric.operations.local>`_ method.
 
@@ -283,9 +297,12 @@ def execute_local(self, command, stdout=True, sudo=False, user=None):
         :type sudo: :class:`bool`
         :param user: If specified, will be the user with which the command will be executed. Defaults to :class:`None`.
         :type user: :class:`str`
+        :param ok_returncodes: List of returncodes that command may return, and that are not treated as an error.
+        :type ok_returncodes: :class:`list` of :clas:`int`.
 
-        :return: The execution result
-        :rtype: :class:`str`
+        :return: The execution result (contents of stdout)
+        :rtype: :class:`str` (a fabric `_AttributeString`, so it has some magic
+            attributes, notably: ``stderr``, ``retcode``, ``failed``.
 
         Example:
         ::
@@ -297,8 +314,11 @@ def provision(self):
                     self.execute_local('ls /', stdout=False, sudo=True)
                     self.execute_local('ls /', stdout=False, user='vip')
         '''
-        with self.__showing_command_output(stdout):
-            return self.__execute_local_command(command, sudo=sudo, user=user)
+        with self.__showing_command_output(stdout), fabric.api.settings(warn_only=True):
+            result =  self.__execute_local_command(command, sudo=sudo, user=user)
+
+        self.__check_returncodes(command, result, ok_returncodes)
+        return result
 
     def __execute_local_command(self, command, sudo=False, user=None):
         if user is not None:

tests/unit/core/test_roles.py

@@ -4,10 +4,12 @@
 from contextlib import contextmanager
 import os
 import tempfile
+from fabric.operations import _AttributeString
 
 from jinja2 import ChoiceLoader, FileSystemLoader
 from mock import MagicMock, patch, call, ANY, Mock, DEFAULT
 from nose.tools import istest
+from provy.core.errors import CommandExecutionError
 
 from provy.core.roles import Role, UsingRole, UpdateData
 from tests.unit.tools.helpers import PROJECT_ROOT, ProvyTestCase
@@ -250,27 +252,69 @@ def can_call_cleanup_safely(self):
     @istest
     def executes_command_with_stdout_and_same_user(self):
         with patch('fabric.api.run') as run:
+            self.__prepare_fabric_patch(run)
             self.role.execute('some command', stdout=True)
 
             run.assert_called_with('some command')
 
     @istest
     def executes_command_with_stdout_and_sudo(self):
         with patch('fabric.api.sudo') as sudo:
+            self.__prepare_fabric_patch(sudo)
             self.role.execute('some command', stdout=True, sudo=True)
 
             sudo.assert_called_with('some command', user=None)
 
+    @istest
+    def executes_command_with_invalid_returncode(self):
+        with patch('fabric.api.run') as run:
+            self.__prepare_fabric_patch(run, return_code=1)
+
+            with self.assertRaises(CommandExecutionError):
+                self.role.execute('some command', stdout=False)
+
+            run.assert_called_with('some command')
+
+    @istest
+    def executes_command_with_ok_returncode(self):
+        with patch('fabric.api.run') as run:
+            self.__prepare_fabric_patch(run, return_code=1)
+
+            self.role.execute('some command', ok_returncodes=(0, 1))
+
+            run.assert_called_with('some command')
+
+    @istest
+    def executes_local_command_with_invalid_returncode(self):
+        with patch('fabric.api.local') as local:
+            self.__prepare_fabric_patch(local, return_code=1)
+
+            with self.assertRaises(CommandExecutionError):
+                self.role.execute_local('some command', stdout=False)
+
+            local.assert_called_with('some command', capture=True)
+
+    @istest
+    def executes_local_command_with_ok_returncode(self):
+        with patch('fabric.api.local') as local:
+            self.__prepare_fabric_patch(local, return_code=1)
+
+            self.role.execute_local('some command', ok_returncodes=(0, 1))
+
+            local.assert_called_with('some command', capture=True)
+
     @istest
     def executes_command_with_stdout_and_another_user(self):
         with patch('fabric.api.sudo') as sudo:
+            self.__prepare_fabric_patch(sudo, 'some_result')
             self.role.execute('some command', stdout=True, user='foo')
 
             sudo.assert_called_with('some command', user='foo')
 
     @istest
     def executes_command_without_stdout_but_same_user(self):
         with patch('fabric.api.run') as run, patch('fabric.api.hide') as hide:
+            self.__prepare_fabric_patch(run)
             self.role.execute('some command', stdout=False)
 
             run.assert_called_with('some command')
@@ -279,6 +323,7 @@ def executes_command_without_stdout_but_same_user(self):
     @istest
     def executes_command_without_stdout_but_sudo(self):
         with patch('fabric.api.sudo') as sudo, patch('fabric.api.hide') as hide:
+            self.__prepare_fabric_patch(sudo)
             self.role.execute('some command', stdout=False, sudo=True)
 
             sudo.assert_called_with('some command', user=None)
@@ -287,54 +332,61 @@ def executes_command_without_stdout_but_sudo(self):
     @istest
     def executes_command_without_stdout_but_another_user(self):
         with patch('fabric.api.sudo') as sudo, patch('fabric.api.hide') as hide:
+            self.__prepare_fabric_patch(sudo)
             self.role.execute('some command', stdout=False, user='foo')
 
             sudo.assert_called_with('some command', user='foo')
             hide.assert_called_with('warnings', 'running', 'stdout', 'stderr')
 
+    def __prepare_fabric_patch(self, patched_obj, str='', return_code=0):
+        result = _AttributeString(str)
+        result.return_code = return_code
+        patched_obj.return_value = result
+        return result
+
     @istest
     def execute_command_check_cd_called_if_cwd_arg(self):
-        with patch('fabric.api.run'):
-            with patch('fabric.api.cd') as cd:
-                self.role.execute("some command", cwd="/some/dir")
+        with patch('fabric.api.run') as run, patch('fabric.api.cd') as cd:
+            self.__prepare_fabric_patch(run)
+            self.role.execute("some command", cwd="/some/dir")
         cd.assert_called_once_with("/some/dir")
 
     @istest
     def execute_command_check_cd_called_if_no_cwd_arg(self):
-        with patch('fabric.api.run'):
-            with patch('fabric.api.cd') as cd:
-                self.role.execute("some command")
+        with patch('fabric.api.run') as run, patch('fabric.api.cd') as cd:
+            self.__prepare_fabric_patch(run)
+            self.role.execute("some command")
         self.assertFalse(cd.called)
 
     @istest
     def executes_a_local_command_with_stdout_and_same_user(self):
         with patch('fabric.api.local') as local:
-            local.return_value = 'some result'
-            self.assertEqual(self.role.execute_local('some command', stdout=True), 'some result')
+            expected_result = self.__prepare_fabric_patch(local, 'some_result')
+            self.assertEqual(self.role.execute_local('some command', stdout=True), expected_result)
 
             local.assert_called_with('some command', capture=True)
 
     @istest
     def executes_a_local_command_with_stdout_and_sudo(self):
         with patch('fabric.api.local') as local:
-            local.return_value = 'some result'
-            self.assertEqual(self.role.execute_local('some command', stdout=True, sudo=True), 'some result')
+            expected_result = self.__prepare_fabric_patch(local, 'some_result')
+            self.assertEqual(self.role.execute_local('some command', stdout=True, sudo=True), expected_result)
 
             local.assert_called_with('sudo some command', capture=True)
 
     @istest
     def executes_a_local_command_with_stdout_and_another_user(self):
         with patch('fabric.api.local') as local:
-            local.return_value = 'some result'
-            self.assertEqual(self.role.execute_local('some command', stdout=True, user='foo'), 'some result')
+            expected_result = self.__prepare_fabric_patch(local, 'some_result')
+            self.assertEqual(self.role.execute_local('some command', stdout=True, user='foo'), expected_result)
 
             local.assert_called_with('sudo -u foo some command', capture=True)
 
     @istest
     def executes_a_local_command_without_stdout_and_another_user(self):
         with patch('fabric.api.local') as local, patch('fabric.api.hide') as hide:
-            local.return_value = 'some result'
-            self.assertEqual(self.role.execute_local('some command', stdout=False, user='foo'), 'some result')
+            expected_result = self.__prepare_fabric_patch(local, 'some_result')
+            self.assertEqual(self.role.execute_local('some command', stdout=False, user='foo'), expected_result)
 
             local.assert_called_with('sudo -u foo some command', capture=True)
             hide.assert_called_with('warnings', 'running', 'stdout', 'stderr')
@@ -343,7 +395,6 @@ def executes_a_local_command_without_stdout_and_another_user(self):
     def executes_a_python_command(self):
         with self.execute_mock() as execute:
             self.role.execute_python('some command', stdout='is stdout?', sudo='is sudo?')
-
             execute.assert_called_with('python -c "some command"', stdout='is stdout?', sudo='is sudo?')
 
     @istest