Merge "Import trycmd and ssh_execute from nova."

openstack/common/processutils.py

@@ -34,6 +34,11 @@ from openstack.common import log as logging
 LOG = logging.getLogger(__name__)
 
 
+class InvalidArgumentError(Exception):
+    def __init__(self, message=None):
+        super(InvalidArgumentError, self).__init__(message)
+
+
 class UnknownArgumentError(Exception):
     def __init__(self, message=None):
         super(UnknownArgumentError, self).__init__(message)
@@ -179,3 +184,64 @@ def execute(*cmd, **kwargs):
             #               call clean something up in between calls, without
             #               it two execute calls in a row hangs the second one
             greenthread.sleep(0)
+
+
+def trycmd(*args, **kwargs):
+    """
+    A wrapper around execute() to more easily handle warnings and errors.
+
+    Returns an (out, err) tuple of strings containing the output of
+    the command's stdout and stderr.  If 'err' is not empty then the
+    command can be considered to have failed.
+
+    :discard_warnings   True | False. Defaults to False. If set to True,
+                        then for succeeding commands, stderr is cleared
+
+    """
+    discard_warnings = kwargs.pop('discard_warnings', False)
+
+    try:
+        out, err = execute(*args, **kwargs)
+        failed = False
+    except ProcessExecutionError, exn:
+        out, err = '', str(exn)
+        failed = True
+
+    if not failed and discard_warnings and err:
+        # Handle commands that output to stderr but otherwise succeed
+        err = ''
+
+    return out, err
+
+
+def ssh_execute(ssh, cmd, process_input=None,
+                addl_env=None, check_exit_code=True):
+    LOG.debug(_('Running cmd (SSH): %s'), cmd)
+    if addl_env:
+        raise InvalidArgumentError(_('Environment not supported over SSH'))
+
+    if process_input:
+        # This is (probably) fixable if we need it...
+        raise InvalidArgumentError(_('process_input not supported over SSH'))
+
+    stdin_stream, stdout_stream, stderr_stream = ssh.exec_command(cmd)
+    channel = stdout_stream.channel
+
+    # NOTE(justinsb): This seems suspicious...
+    # ...other SSH clients have buffering issues with this approach
+    stdout = stdout_stream.read()
+    stderr = stderr_stream.read()
+    stdin_stream.close()
+
+    exit_status = channel.recv_exit_status()
+
+    # exit_status == -1 if no exit code was returned
+    if exit_status != -1:
+        LOG.debug(_('Result was %s') % exit_status)
+        if check_exit_code and exit_status != 0:
+            raise ProcessExecutionError(exit_code=exit_status,
+                                        stdout=stdout,
+                                        stderr=stderr,
+                                        cmd=cmd)
+
+    return (stdout, stderr)

tests/unit/test_processutils.py

@@ -17,7 +17,9 @@
 
 from __future__ import print_function
 
+import fixtures
 import os
+import StringIO
 import tempfile
 
 from openstack.common import processutils
@@ -164,3 +166,86 @@ grep foo
         finally:
             os.unlink(tmpfilename)
             os.unlink(tmpfilename2)
+
+
+def fake_execute(*cmd, **kwargs):
+    return 'stdout', 'stderr'
+
+
+def fake_execute_raises(*cmd, **kwargs):
+    raise processutils.ProcessExecutionError(exit_code=42,
+                                             stdout='stdout',
+                                             stderr='stderr',
+                                             cmd=['this', 'is', 'a',
+                                                  'command'])
+
+
+class TryCmdTestCase(utils.BaseTestCase):
+    def test_keep_warnings(self):
+        self.useFixture(fixtures.MonkeyPatch(
+            'openstack.common.processutils.execute', fake_execute))
+        o, e = processutils.trycmd('this is a command'.split(' '))
+        self.assertNotEqual('', o)
+        self.assertNotEqual('', e)
+
+    def test_keep_warnings_from_raise(self):
+        self.useFixture(fixtures.MonkeyPatch(
+            'openstack.common.processutils.execute', fake_execute_raises))
+        o, e = processutils.trycmd('this is a command'.split(' '),
+                                   discard_warnings=True)
+        self.assertNotEqual(None, o)
+        self.assertNotEqual('', e)
+
+    def test_discard_warnings(self):
+        self.useFixture(fixtures.MonkeyPatch(
+            'openstack.common.processutils.execute', fake_execute))
+        o, e = processutils.trycmd('this is a command'.split(' '),
+                                   discard_warnings=True)
+        self.assertNotEqual(None, o)
+        self.assertEqual('', e)
+
+
+class FakeSshChannel(object):
+    def __init__(self, rc):
+        self.rc = rc
+
+    def recv_exit_status(self):
+        return self.rc
+
+
+class FakeSshStream(StringIO.StringIO):
+    def setup_channel(self, rc):
+        self.channel = FakeSshChannel(rc)
+
+
+class FakeSshConnection(object):
+    def __init__(self, rc):
+        self.rc = rc
+
+    def exec_command(self, cmd):
+        stdout = FakeSshStream('stdout')
+        stdout.setup_channel(self.rc)
+        return (StringIO.StringIO(),
+                stdout,
+                StringIO.StringIO('stderr'))
+
+
+class SshExecuteTestCase(utils.BaseTestCase):
+    def test_invalid_addl_env(self):
+        self.assertRaises(processutils.InvalidArgumentError,
+                          processutils.ssh_execute,
+                          None, 'ls', addl_env='important')
+
+    def test_invalid_process_input(self):
+        self.assertRaises(processutils.InvalidArgumentError,
+                          processutils.ssh_execute,
+                          None, 'ls', process_input='important')
+
+    def test_works(self):
+        o, e = processutils.ssh_execute(FakeSshConnection(0), 'ls')
+        self.assertEqual('stdout', o)
+        self.assertEqual('stderr', e)
+
+    def test_fails(self):
+        self.assertRaises(processutils.ProcessExecutionError,
+                          processutils.ssh_execute, FakeSshConnection(1), 'ls')