Make cmd starting async. Resolves #56. Updated tests and documentation.

Author: Dan

pssh/pssh_client.py

@@ -163,6 +163,15 @@ def __init__(self, hosts,
         >>> import paramiko
         >>> client_key = paramiko.RSAKey.from_private_key_file('user.key')
         >>> client = ParallelSSHClient(['myhost1', 'myhost2'], pkey=client_key)
+
+        **Example with expression as host list**
+        
+        Any type of iterator may be used as host list, including generator and
+        list comprehension expressions.
+        
+        >>> hosts = ['dc1.myhost1', 'dc2.myhost2']
+        >>> client = ParallelSSHClient([h for h in hosts if h.find('dc1')])
+        >>> client.run_command(<..>)
         
         .. note ::
         
@@ -184,7 +193,7 @@ def __init__(self, hosts,
           
           Connection is terminated.
         """
-        self.pool_size = len(hosts) if len(hosts) < pool_size else pool_size
+        self.pool_size = pool_size
         self.pool = gevent.pool.Pool(size=self.pool_size)
         self.hosts = hosts
         self.user = user
@@ -202,15 +211,21 @@ def __init__(self, hosts,
     def run_command(self, *args, **kwargs):
         """Run command on all hosts in parallel, honoring self.pool_size,
         and return output buffers.
-
+        
         This function will block until all commands have **started** and
-        then return immediately. Any connection and/or authentication exceptions
-        will be raised here and need catching.
-
+        then return immediately.
+        
+        Any connection and/or authentication exceptions will be raised here
+        and need catching _unless_ `run_command` is called with
+        `stop_on_errors=False`.
+        
         :param args: Positional arguments for command
         :type args: tuple
         :param sudo: (Optional) Run with sudo. Defaults to False
         :type sudo: bool
+        :param user: (Optional) User to run command as. Requires sudo access \
+        for that user from the logged in user account.
+        :type user: str
         :param stop_on_errors: (Optional) Raise exception on errors running command. \
         Defaults to True. With stop_on_errors set to False, exceptions are instead \
         added to output of `run_command`. See example usage below.
@@ -235,18 +250,20 @@ def run_command(self, *args, **kwargs):
         
         >>> for host in output:
         >>>     for line in output[host]['stdout']: print line
-
+        
         *Get exit codes after command has finished*
-
+        
         >>> client.get_exit_codes(output)
         >>> for host in output:
         >>> ... print output[host]['exit_code']
         0
         0
         
-        *Wait for completion, no stdout printing*
+        *Wait for completion, no stdout/stderr*
         
         >>> client.join(output)
+        >>> print output[host]['exit_code']
+        0
         
         *Run with sudo*
         
@@ -370,33 +387,40 @@ def get_output(self, cmd, output):
             try:
                 host = ex.args[1]
             except IndexError:
-                logger.error("Got exception with no host argument - cannot update output data with %s", ex)
+                logger.error("Got exception with no host argument - "
+                             "cannot update output data with %s", ex)
                 raise ex
             self._update_host_output(output, host, None, None, None, None, cmd,
                                      exception=ex)
             raise ex
         self._update_host_output(output, host, self._get_exit_code(channel),
                                  channel, stdout, stderr, cmd)
 
-    def _update_host_output(self, output, host, exit_code, channel, stdout, stderr, cmd,
-                            exception=None):
+    def _update_host_output(self, output, host, exit_code, channel, stdout,
+                            stderr, cmd, exception=None):
         """Update host output with given data"""
         if host in output:
             new_host = "_".join([host,
                                  ''.join(random.choice(
                                      string.ascii_lowercase + string.digits)
                                      for _ in xrange(8))])
-            logger.warning("Already have output for host %s - changing host key for %s to %s",
-                           host, host, new_host)
+            logger.warning("Already have output for host %s - changing host "
+                           "key for %s to %s", host, host, new_host)
             host = new_host
         output.setdefault(host, {})
         output[host].update({'exit_code' : exit_code,
                              'channel' : channel,
-                             'stdout' : stdout,
-                             'stderr' : stderr,
+                             'stdout' : self._read_buff_ex_code(stdout, output),
+                             'stderr' : self._read_buff_ex_code(stderr, output),
                              'cmd' : cmd,
                              'exception' : exception,})
-    
+
+    def _read_buff_ex_code(self, _buffer, output):
+        if _buffer:
+            for line in _buffer:
+                yield line
+        self.get_exit_codes(output)
+
     def join(self, output):
         """Block until all remote commands in output have finished
         and retrieve exit codes"""

pssh/ssh_client.py

@@ -229,9 +229,7 @@ def exec_command(self, command, sudo=False, user=None, **kwargs):
         logger.debug("Running command %s on %s", command, self.host)
         channel.exec_command(command, **kwargs)
         logger.debug("Command started")
-        while not (channel.recv_ready() or channel.closed or
-                   channel.exit_status_ready()):
-            gevent.sleep(.2)
+        gevent.sleep(0)
         return channel, self.host, stdout, stderr
 
     def _read_output_buffer(self, output_buffer, prefix=''):

tests/test_pssh_client.py

@@ -74,8 +74,8 @@ def test_pssh_client_exec_command(self):
     def test_pssh_client_no_stdout_non_zero_exit_code(self):
         output = self.client.run_command('exit 1')
         expected_exit_code = 1
+        self.client.join(output)
         exit_code = output[self.host]['exit_code']
-        self.client.pool.join()
         self.assertEqual(expected_exit_code, exit_code,
                          msg="Got unexpected exit code - %s, expected %s" %
                          (exit_code,
@@ -114,9 +114,9 @@ def test_pssh_client_run_command_get_output(self):
         expected_exit_code = 0
         expected_stdout = [self.fake_resp]
         expected_stderr = []
-        exit_code = output[self.host]['exit_code']
         stdout = list(output[self.host]['stdout'])
         stderr = list(output[self.host]['stderr'])
+        exit_code = output[self.host]['exit_code']
         self.assertEqual(expected_exit_code, exit_code,
                          msg="Got unexpected exit code - %s, expected %s" %
                          (exit_code,
@@ -141,9 +141,9 @@ def test_pssh_client_run_command_get_output_explicit(self):
         expected_exit_code = 0
         expected_stdout = [self.fake_resp]
         expected_stderr = []
-        exit_code = output[self.host]['exit_code']
         stdout = list(output[self.host]['stdout'])
         stderr = list(output[self.host]['stderr'])
+        exit_code = output[self.host]['exit_code']
         self.assertEqual(expected_exit_code, exit_code,
                          msg="Got unexpected exit code - %s, expected %s" %
                          (exit_code,
@@ -378,22 +378,25 @@ def test_pssh_client_directory(self):
         shutil.rmtree(remote_test_path)
 
     def test_pssh_pool_size(self):
-        """Test pool size logic"""
+        """Test setting pool size to non default values"""
         hosts = ['host-%01d' % d for d in xrange(5)]
-        client = ParallelSSHClient(hosts)
-        expected, actual = len(hosts), client.pool.size
+        pool_size = 2
+        client = ParallelSSHClient(hosts, pool_size=pool_size)
+        expected, actual = pool_size, client.pool.size
         self.assertEqual(expected, actual,
                          msg="Expected pool size to be %s, got %s" % (
                              expected, actual,))
         hosts = ['host-%01d' % d for d in xrange(15)]
-        client = ParallelSSHClient(hosts)
+        pool_size = 5
+        client = ParallelSSHClient(hosts, pool_size=pool_size)
         expected, actual = client.pool_size, client.pool.size
         self.assertEqual(expected, actual,
                          msg="Expected pool size to be %s, got %s" % (
                              expected, actual,))
         hosts = ['host-%01d' % d for d in xrange(15)]
-        client = ParallelSSHClient(hosts, pool_size=len(hosts)+5)
-        expected, actual = len(hosts), client.pool.size
+        pool_size = len(hosts)+5
+        client = ParallelSSHClient(hosts, pool_size=pool_size)
+        expected, actual = pool_size, client.pool.size
         self.assertEqual(expected, actual,
                          msg="Expected pool size to be %s, got %s" % (
                              expected, actual,))

tests/test_ssh_client.py

@@ -205,7 +205,6 @@ def test_ssh_agent_authentication(self):
         client = SSHClient(self.host, port=self.listen_port,
                            agent=agent)
         channel, host, stdout, stderr = client.exec_command(self.fake_cmd)
-        channel.close()
         output = list(stdout)
         stderr = list(stderr)
         expected = [self.fake_resp]