Added per-host command arguments and tests. Resolves #43

Author: Dan

.travis.yml

@@ -15,9 +15,6 @@ script:
   - python setup.py nosetests --with-coverage --cover-package=pssh
   # If using py3, run 2to3 on embedded server and tests and run nosetests for new test dir
   - python -c 'import sys; sys.version_info >= (3,) and sys.exit(1)' || eval "2to3 -nw embedded_server/*.py && 2to3 tests/*.py -o tests3 -nw && cp tests/test_client_private_key* tests3/ && python setup.py nosetests -w tests3 --with-coverage --cover-package=pssh"
-notifications:
-  email:
-    on_failure: change
 after_success:
   - coveralls
 deploy:

pssh/exceptions.py

@@ -37,3 +37,8 @@ class AuthenticationException(Exception):
 class SSHException(Exception):
     """Raised on SSHException error - error authenticating with SSH server"""
     pass
+
+
+class HostArgumentException(Exception):
+    """Raised on errors with per-host command arguments"""
+    pass

pssh/pssh_client.py

@@ -30,6 +30,8 @@
 import warnings
 import string
 import random
+
+from .exceptions import HostArgumentException
 from .constants import DEFAULT_RETRIES
 from .ssh_client import SSHClient
 
@@ -348,13 +350,21 @@ def run_command(self, *args, **kwargs):
         :param use_shell: (Optional) Run command with or without shell. Defaults \
         to True - use shell defined in user login to run command string
         :type use_shell: bool
+        :param host_args: (Optional) Format command string with per-host \
+        arguments in ``host_args``. ``host_args`` length must equal length of \
+        host list - :mod:`pssh.exceptions.HostArgumentException` is raised \
+        otherwise
+        :type host_args: tuple or list
         :rtype: Dictionary with host as key as per \
           :mod:`pssh.pssh_client.ParallelSSHClient.get_output`
         
         :raises: :mod:`pssh.exceptions.AuthenticationException` on authentication error
         :raises: :mod:`pssh.exceptions.UnknownHostException` on DNS resolution error
         :raises: :mod:`pssh.exceptions.ConnectionErrorException` on error connecting
         :raises: :mod:`pssh.exceptions.SSHException` on other undefined SSH errors
+        :raises: :mod:`pssh.exceptions.HostArgumentException` on number of host \
+        arguments not equal to number of hosts
+        :raises: `TypeError` on not enough host arguments for cmd string format
         
         **Example Usage**
         
@@ -364,14 +374,15 @@ def run_command(self, *args, **kwargs):
         
           output = client.run_command('ls -ltrh')
         
-        *print stdout for each command*
+        **Print stdout for each command**
         
         .. code-block:: python
         
           for host in output:
-              for line in output[host]['stdout']: print(line)
+              for line in output[host]['stdout']:
+                  print(line)
         
-        *Get exit codes after command has finished*
+        **Get exit codes after command has finished**
 
         .. code-block:: python
         
@@ -391,7 +402,8 @@ def run_command(self, *args, **kwargs):
           for line in output[host]['stdout']:
               print(line)
         
-        *Run with sudo*
+        **Run with sudo**
+        
 
         .. code-block:: python
         
@@ -410,8 +422,45 @@ def run_command(self, *args, **kwargs):
           for host in output:
               stdout = list(output[host]['stdout'])
               print("Complete stdout for host %s is %s" % (host, stdout,))
+
+        **Command with per-host arguments**
+
+        ``host_args`` keyword parameter can be used to provide string formatting
+        arguments to use to format the command string.
+
+        ``host_args`` should be at least as many as number of hosts.
+
+        Any string format specification characters may be used in command string.
+
+        *Examples*
+        
+        .. code-block:: python
+
+          # Tuple
+          # 
+          # First host in hosts list will use cmd 'host1_cmd',
+          # second host 'host2_cmd' and so on
+          output = client.run_command('%s', host_args=('host1_cmd',
+                                                       'host2_cmd',
+                                                       'host3_cmd',))
+          
+          # Multiple arguments
+          #
+          output = client.run_command('%s %s',
+                                      host_args=(('host1_cmd1', 'host1_cmd2'),
+                                                 ('host2_cmd1', 'host2_cmd2'),
+                                                 ('host3_cmd1', 'host3_cmd2'),))
+          
+          # List of dict
+          # 
+          # Fist host in host list will use cmd 'host-index-0',
+          # second host 'host-index-1' and so on
+          output = client.run_command(
+            '%(cmd)s', host_args=[{'cmd': 'host-index-%s' % (i,))
+                                  for i in range(len(client.hosts))]
         
         **Expression as host list**
+
         
         Any type of iterator may be used as host list, including generator and
         list comprehension expressions.
@@ -477,7 +526,7 @@ def run_command(self, *args, **kwargs):
           Started 10 commands in 0:00:00.428629
           All commands finished in 0:00:05.014757
         
-        **Output dictionary**
+        *Output dictionary*
         
         ::
         
@@ -522,8 +571,20 @@ def run_command(self, *args, **kwargs):
         
         """
         stop_on_errors = kwargs.pop('stop_on_errors', True)
-        cmds = [self.pool.spawn(self._exec_command, host, *args, **kwargs)
-                for host in self.hosts]
+        host_args = kwargs.pop('host_args', None)
+        if host_args:
+            try:
+                cmds = [self.pool.spawn(self._exec_command, host,
+                                        args[0] % host_args[host_i],
+                                        *args[1:], **kwargs)
+                        for host_i, host in enumerate(self.hosts)]
+            except IndexError:
+                raise HostArgumentException(
+                    "Number of host arguments provided does not match "
+                    "number of hosts ")
+        else:
+            cmds = [self.pool.spawn(self._exec_command, host, *args, **kwargs)
+                    for host in self.hosts]
         output = {}
         for cmd in cmds:
             try:

tests/test_pssh_client.py

@@ -19,25 +19,28 @@
 
 """Unittests for :mod:`pssh.ParallelSSHClient` class"""
 
+
 import unittest
-from pssh import ParallelSSHClient, UnknownHostException, \
-     AuthenticationException, ConnectionErrorException, SSHException, logger as pssh_logger
-from pssh.utils import load_private_key
-from embedded_server.embedded_server import start_server, make_socket, \
-     logger as server_logger, paramiko_logger
-from embedded_server.fake_agent import FakeAgent
 import random
 import logging
-import gevent
-import paramiko
 import os
 import warnings
 import shutil
 import sys
 
+import gevent
+from pssh import ParallelSSHClient, UnknownHostException, \
+     AuthenticationException, ConnectionErrorException, SSHException, \
+     logger as pssh_logger
+from pssh.exceptions import HostArgumentException
+from pssh.utils import load_private_key
+from embedded_server.embedded_server import start_server, make_socket, \
+     logger as server_logger, paramiko_logger
+from embedded_server.fake_agent import FakeAgent
+from paramiko import RSAKey
 
 PKEY_FILENAME = os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key'])
-USER_KEY = paramiko.RSAKey.from_private_key_file(PKEY_FILENAME)
+USER_KEY = RSAKey.from_private_key_file(PKEY_FILENAME)
 
 server_logger.setLevel(logging.DEBUG)
 pssh_logger.setLevel(logging.DEBUG)
@@ -228,7 +231,7 @@ def test_pssh_client_ssh_exception(self):
         client = ParallelSSHClient([self.host],
                                    user='fakey', password='fakey',
                                    port=listen_port,
-                                   pkey=paramiko.RSAKey.generate(1024),
+                                   pkey=RSAKey.generate(1024),
                                    )
         self.assertRaises(SSHException, client.run_command, self.fake_cmd)
         del client
@@ -710,7 +713,7 @@ def test_ssh_exception(self):
         hosts = [host]
         client = ParallelSSHClient(hosts, port=port,
                                    user='fakey', password='fakey',
-                                   pkey=paramiko.RSAKey.generate(1024))
+                                   pkey=RSAKey.generate(1024))
         output = client.run_command(self.fake_cmd, stop_on_errors=False)
         gevent.sleep(1)
         client.pool.join()
@@ -826,7 +829,6 @@ def test_pssh_client_override_allow_agent_authentication(self):
         expected_exit_code = 0
         expected_stdout = [self.fake_resp]
         expected_stderr = []
-
         stdout = list(output[self.host]['stdout'])
         stderr = list(output[self.host]['stderr'])
         exit_code = output[self.host]['exit_code']
@@ -848,5 +850,69 @@ def test_get_exit_codes_bad_output(self):
         self.assertFalse(self.client.get_exit_codes({}))
         self.assertFalse(self.client.get_exit_code({}))
 
+    def test_per_host_tuple_args(self):
+        server2_socket = make_socket('127.0.0.2', port=self.listen_port)
+        server2_port = server2_socket.getsockname()[1]
+        server2 = start_server(server2_socket)
+        server3_socket = make_socket('127.0.0.3', port=self.listen_port)
+        server3_port = server3_socket.getsockname()[1]
+        server3 = start_server(server3_socket)
+        hosts = [self.host, '127.0.0.2', '127.0.0.3']
+        host_args = ('arg1', 'arg2', 'arg3')
+        cmd = 'echo %s'
+        client = ParallelSSHClient(hosts, port=self.listen_port,
+                                   pkey=self.user_key)
+        output = client.run_command(cmd, host_args=host_args)
+        for i, host in enumerate(hosts):
+            expected = [host_args[i]]
+            stdout = list(output[host]['stdout'])
+            self.assertEqual(expected, stdout)
+            self.assertTrue(output[host]['exit_code'] == 0)
+        host_args = (('arg1', 'arg2'), ('arg3', 'arg4'), ('arg5', 'arg6'),)
+        cmd = 'echo %s %s'
+        output = client.run_command(cmd, host_args=host_args)
+        for i, host in enumerate(hosts):
+            expected = ["%s %s" % host_args[i]]
+            stdout = list(output[host]['stdout'])
+            self.assertEqual(expected, stdout)
+            self.assertTrue(output[host]['exit_code'] == 0)
+        self.assertRaises(HostArgumentException, client.run_command,
+                          cmd, host_args=[host_args[0]])
+
+    def test_per_host_dict_args(self):
+        server2_socket = make_socket('127.0.0.2', port=self.listen_port)
+        server2_port = server2_socket.getsockname()[1]
+        server2 = start_server(server2_socket)
+        server3_socket = make_socket('127.0.0.3', port=self.listen_port)
+        server3_port = server3_socket.getsockname()[1]
+        server3 = start_server(server3_socket)
+        hosts = [self.host, '127.0.0.2', '127.0.0.3']
+        hosts_gen = (h for h in hosts)
+        host_args = [dict(zip(('host_arg1', 'host_arg2',),
+                              ('arg1-%s' % (i,), 'arg2-%s' % (i,),)))
+                     for i, _ in enumerate(hosts)]
+        cmd = 'echo %(host_arg1)s %(host_arg2)s'
+        client = ParallelSSHClient(hosts, port=self.listen_port,
+                                   pkey=self.user_key)
+        output = client.run_command(cmd, host_args=host_args)
+        for i, host in enumerate(hosts):
+            expected = ["%(host_arg1)s %(host_arg2)s" % host_args[i]]
+            stdout = list(output[host]['stdout'])
+            self.assertEqual(expected, stdout)
+            self.assertTrue(output[host]['exit_code'] == 0)
+        self.assertRaises(HostArgumentException, client.run_command,
+                          cmd, host_args=[host_args[0]])
+        # Host list generator should work also
+        client.hosts = hosts_gen
+        output = client.run_command(cmd, host_args=host_args)
+        for i, host in enumerate(hosts):
+            expected = ["%(host_arg1)s %(host_arg2)s" % host_args[i]]
+            stdout = list(output[host]['stdout'])
+            self.assertEqual(expected, stdout)
+            self.assertTrue(output[host]['exit_code'] == 0)
+        client.hosts = (h for h in hosts)
+        self.assertRaises(HostArgumentException, client.run_command,
+                          cmd, host_args=[host_args[0]])
+
 if __name__ == '__main__':
     unittest.main()