Updated part failure test to also check for presence of (correct) exception in output for failed host. Added enable host logger function to pssh module

Author: Dan

pssh.py

@@ -64,6 +64,20 @@ class SSHException(Exception):
     pass
 
 
+def enable_host_logger():
+    """Enables host logger for logging stdout from remote servers as it
+    becomes available
+    """
+    if logging.StreamHandler in [type(h) for h in host_logger.handlers]:
+        logger.warning("Host logger already has a StreamHandler attached")
+        return
+    handler = logging.StreamHandler()
+    host_log_format = logging.Formatter('%(message)s')
+    handler.setFormatter(host_log_format)
+    host_logger.addHandler(handler)
+    host_logger.setLevel(logging.INFO)
+    
+
 class SSHClient(object):
     """Wrapper class over paramiko.SSHClient with sane defaults
     Honours ~/.ssh/config and /etc/ssh/ssh_config entries for host username \
@@ -540,22 +554,24 @@ def get_output(self, cmd, output):
         
         :param cmd: Command to get output from
         :type cmd: :mod:`gevent.Greenlet`
-        :rtype: Dictionary with host as key as in:
+        :param output: Dictionary containing output to be updated with output
+        from cmd
+        :type output: dict
+        :rtype: None
 
+        `output` parameter is modified in-place and has the following structure
+        
         ::
         
           {'myhost1': {'exit_code': exit code if ready else None,
                        'channel' : SSH channel of command,
                        'stdout'  : <iterable>,
                        'stderr'  : <iterable>,
-                       'cmd'     : <greenlet>}}
+                       'cmd'     : <greenlet>,
+                       'exception' : <exception object if applicable>}}
         
         Stdout and stderr are also logged via the logger named ``host_logger``
-        which is enabled by default.
-        ``host_logger`` output can be disabled by removing its handler.
-        
-        >>> logger = logging.getLogger('pssh.host_logger')
-        >>> for handler in logger.handlers: logger.removeHandler(handler)
+        which can be enabled by calling ``enable_host_logger``
         
         **Example usage**:
         

tests/test_pssh_client.py

@@ -194,10 +194,7 @@ def test_pssh_client_hosts_list_part_failure(self):
         host in the host list has a failure"""
         server2_socket = make_socket('127.0.0.2', port=self.listen_port)
         server2_port = server2_socket.getsockname()[1]
-        server1 = start_server({ self.fake_cmd : self.fake_resp },
-                               self.listen_socket, fail_auth=True)
-        server2 = start_server({ self.fake_cmd : self.fake_resp },
-                               server2_socket)
+        server2 = start_server(server2_socket, fail_auth=True)
         hosts = ['127.0.0.1', '127.0.0.2']
         client = ParallelSSHClient(hosts,
                                    port=self.listen_port,
@@ -206,13 +203,21 @@ def test_pssh_client_hosts_list_part_failure(self):
         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,
                         msg="Successful host does not exist in output - output is %s" % (output,))
+        self.assertTrue(hosts[1] in output,
+                        msg="Failed host does not exist in output - output is %s" % (output,))
+        self.assertTrue('exception' in output[hosts[1]],
+                        msg="Failed host %s has no exception in output - %s" % (hosts[1], output,))
+        try:
+            raise output[hosts[1]]['exception']
+        except AuthenticationException:
+            pass
+        else:
+            raise Exception("Expected AuthenticationException, got %s instead" % (
+                output[hosts[1]]['exception'],))
         del client
-        server1.kill()
         server2.kill()
-
+    
     def test_pssh_client_ssh_exception(self):
         listen_socket = make_socket('127.0.0.1')
         listen_port = listen_socket.getsockname()[1]