Updated get_output and exec_command for changed behaviour. Added host as argument to exceptions so that output can be updated correctly for the host the exception occured for. Updated test_pssh_client_run_command_get_output_explicit for changed get_output interface - Resolves #32

Author: Dan

pssh.py

@@ -210,7 +210,7 @@ def _connect(self, client, host, port, sock=None, retries=1):
         # of SSH failure
         except paramiko.SSHException, ex:
             logger.error("General SSH error - %s", ex)
-            raise SSHException(ex.message)
+            raise SSHException(ex.message, host)
 
     def exec_command(self, command, sudo=False, user=None, **kwargs):
         """Wrapper to :mod:`paramiko.SSHClient.exec_command`
@@ -498,15 +498,13 @@ def run_command(self, *args, **kwargs):
                        'cmd'     : <greenlet>}}
 
         """
-        stop_on_errors = kwargs['stop_on_errors'] \
-          if 'stop_on_errors' in kwargs else True
-        del kwargs['stop_on_errors']
+        stop_on_errors = kwargs.pop('stop_on_errors', True)
         cmds = [self.pool.spawn(self._exec_command, host, *args, **kwargs)
                 for host in self.hosts]
         output = {}
         for cmd in cmds:
             try:
-                output = self.get_output(cmd)
+                self.get_output(cmd, output)
             except Exception, ex:
                 if stop_on_errors:
                     raise ex
@@ -541,12 +539,11 @@ def _exec_command(self, host, *args, **kwargs):
                                                 proxy_port=self.proxy_port)
         return self.host_clients[host].exec_command(*args, **kwargs)
 
-    def get_output(self, cmd):
+    def get_output(self, cmd, output):
         """Get output from command.
         
-        :param commands: (Optional) Override commands to get output from.\
-        Uses running commands in pool if not given.
-        :type commands: :mod:`gevent.Greenlet`
+        :param cmd: Command to get output from
+        :type cmd: :mod:`gevent.Greenlet`
         :rtype: Dictionary with host as key as in:
 
         ::
@@ -573,10 +570,15 @@ def get_output(self, cmd):
         >>> self.get_exit_code(output[host])
         0
         """
-        output = {}
         try:
             (channel, host, stdout, stderr) = cmd.get()
         except Exception, ex:
+            try:
+                host = ex.args[1]
+            except IndexError:
+                logger.error("Got exception with no host argument - cannot update output data with %s", ex)
+                raise ex
+            output.setdefault(host, {})
             output[host].update({'exit_code' : None,
                                  'channel' : None,
                                  'stdout' : None,
@@ -590,7 +592,6 @@ def get_output(self, cmd):
                              'stdout' : stdout,
                              'stderr' : stderr,
                              'cmd' : cmd, })
-        return output
 
     def get_exit_code(self, host_output):
         """Get exit code from host output if available

tests/test_pssh_client.py

@@ -137,7 +137,9 @@ def test_pssh_client_run_command_get_output_explicit(self):
                                    pkey=self.user_key)
         out = client.run_command(self.fake_cmd)
         cmds = [cmd for host in out for cmd in [out[host]['cmd']]]
-        output = client.get_output(commands=cmds)
+        output = {}
+        for cmd in cmds:
+            client.get_output(cmd, output)
         expected_exit_code = 0
         expected_stdout = [self.fake_resp]
         expected_stderr = []
@@ -204,13 +206,8 @@ def test_pssh_client_hosts_list_part_failure(self):
                                    port=self.listen_port,
                                    pkey=self.user_key,
                                    )
-        output = {}
-        try:
-            client.run_command(self.fake_cmd,
-                               output=output,
-                               stop_on_errors=False)
-        except AuthenticationException:
-            pass
+        output = client.run_command(self.fake_cmd,
+                                    stop_on_errors=False)
         self.assertTrue(hosts[0] in output,
                         msg="Failed host does not exist in output - output is %s" % (output,))
         self.assertTrue(hosts[1] in output,