Updated run_command to use named arguments. Added environment keyword argument to run_command and unittest. Added unittests for run_command parameters. Updated docstrings

Author: Dan

pssh/pssh_client.py

@@ -335,7 +335,9 @@ def __init__(self, hosts,
         self.host_config = host_config if host_config else {}
         self.channel_timeout = channel_timeout
 
-    def run_command(self, *args, **kwargs):
+    def run_command(self, command, sudo=False, user=None, stop_on_errors=True,
+                    shell=None, use_shell=True, use_pty=True, host_args=None,
+                    environment=None):
         """Run command on all hosts in parallel, honoring self.pool_size,
         and return output buffers.
 
@@ -351,8 +353,8 @@ def run_command(self, *args, **kwargs):
         ``stop_on_errors=False`` in which case exceptions are added to host
         output instead.
 
-        :param args: Positional arguments for command
-        :type args: tuple
+        :param command: Command to run
+        :type command: str
         :param sudo: (Optional) Run with sudo. Defaults to False
         :type sudo: bool
         :param user: (Optional) User to run command as. Requires sudo access \
@@ -382,6 +384,11 @@ def run_command(self, *args, **kwargs):
         host list - :py:class:`pssh.exceptions.HostArgumentException` is raised \
         otherwise
         :type host_args: tuple or list
+        :param environment: (Optional) Environment variables to be exposed to \
+        command to be run. This requires that ``AcceptEnv`` server setting \
+        is enabled - variables will silently not be set otherwise
+        :type environment: dict
+
         :rtype: Dictionary with host as key and \
         :py:class:`pssh.output.HostOutput` as value as per \
         :py:func:`pssh.pssh_client.ParallelSSHClient.get_output`
@@ -617,22 +624,25 @@ def run_command(self, *args, **kwargs):
           writing to stdin
         
         """
-        stop_on_errors = kwargs.pop('stop_on_errors', True)
-        host_args = kwargs.pop('host_args', None)
         output = {}
         if host_args:
             try:
                 cmds = [self.pool.spawn(self._exec_command, host,
-                                        args[0] % host_args[host_i],
-                                        *args[1:], **kwargs)
+                                        command % host_args[host_i],
+                                        sudo=sudo, user=user, shell=shell,
+                                        use_shell=use_shell, use_pty=use_pty,
+                                        environment=environment)
                         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)
+                self._exec_command, host, command,
+                sudo=sudo, user=user, shell=shell,
+                use_shell=use_shell, use_pty=use_pty,
+                environment=environment)
                 for host in self.hosts]
         for cmd in cmds:
             try:
@@ -648,11 +658,14 @@ def _get_host_config_values(self, host):
         _password = self.host_config.get(host, {}).get('password', self.password)
         _pkey = self.host_config.get(host, {}).get('private_key', self.pkey)
         return _user, _port, _password, _pkey
-    
-    def _exec_command(self, host, *args, **kwargs):
+
+    def _exec_command(self, host, command, sudo=False, user=None,
+                      shell=None, use_shell=True, use_pty=True,
+                      environment=None):
         """Make SSHClient, run command on host"""
         if not host in self.host_clients or not self.host_clients[host]:
             _user, _port, _password, _pkey = self._get_host_config_values(host)
+            _user = user if user else _user
             self.host_clients[host] = SSHClient(host, user=_user,
                                                 password=_password,
                                                 port=_port, pkey=_pkey,
@@ -667,21 +680,23 @@ def _exec_command(self, host, *args, **kwargs):
                                                 allow_agent=self.allow_agent,
                                                 agent=self.agent,
                                                 channel_timeout=self.channel_timeout)
-        return self.host_clients[host].exec_command(*args, **kwargs)
+        return self.host_clients[host].exec_command(
+            command, sudo=sudo, user=user, shell=shell,
+            use_shell=use_shell, use_pty=use_pty, environment=environment)
 
     def get_output(self, cmd, output):
         """Get output from command.
-        
+
         :param cmd: Command to get output from
         :type cmd: :py:class:`gevent.Greenlet`
         :param output: Dictionary containing \
         :py:class:`pssh.output.HostOutput` values to be updated with output \
         from cmd
         :type output: dict
         :rtype: None
-        
+
         `output` parameter is modified in-place and has the following structure
-        
+
         ::
         
           {'myhost1':
@@ -705,7 +720,8 @@ def get_output(self, cmd, output):
               for line in output[host].stdout:
                   print(line)
           <stdout>
-          # Get exit code after command has finished
+          # Get exit code for a particular host's output after command
+          # has finished
           self.get_exit_code(output[host])
           0
         
@@ -763,7 +779,7 @@ def finished(self, output):
         :rtype: bool
         """
         for host in output:
-            chan = output[host]['channel']
+            chan = output[host].channel
             if chan is not None and not chan.closed:
                 return False
         return True
@@ -787,7 +803,7 @@ def get_exit_code(self, host_output):
         if not 'channel' in host_output:
             logger.error("%s does not look like host output..", host_output,)
             return
-        channel = host_output['channel']
+        channel = host_output.channel
         return self._get_exit_code(channel)
 
     def _get_exit_code(self, channel):

pssh/ssh_client.py

@@ -203,7 +203,7 @@ def _connect(self, client, host, port, sock=None, retries=1,
     def exec_command(self, command, sudo=False, user=None,
                      shell=None,
                      use_shell=True, use_pty=True,
-                     **kwargs):
+                     environment=None):
         """Wrapper to :py:func:`paramiko.SSHClient.exec_command`
         
         Opens a new SSH session with a new pty and runs command before yielding 
@@ -241,6 +241,8 @@ def exec_command(self, command, sudo=False, user=None,
             channel.get_pty()
         if self.channel_timeout:
             channel.settimeout(self.channel_timeout)
+        if environment:
+            channel.update_environment(environment)
         stdout, stderr, stdin = channel.makefile('rb'), channel.makefile_stderr('rb'), \
           channel.makefile('wb')
         for _char in ['\\', '"', '$', '`']:
@@ -256,7 +258,7 @@ def exec_command(self, command, sudo=False, user=None,
         else:
             _command += '"%s"' % (command,)
         logger.debug("Running parsed command %s on %s", _command, self.host)
-        channel.exec_command(_command, **kwargs)
+        channel.exec_command(_command)
         logger.debug("Command started")
         sleep(0)
         return channel, self.host, stdout, stderr, stdin

tests/test_pssh_client.py

@@ -976,5 +976,35 @@ def test_output_attributes(self):
         self.assertTrue(hasattr(output[self.host], 'exception'))
         self.assertTrue(hasattr(output[self.host], 'exit_code'))
 
+    def test_run_command_user_sudo(self):
+        user = 'cmd_user'
+        output = self.client.run_command(self.fake_cmd, user=user)
+        self.client.join(output)
+        stderr = list(output[self.host].stderr)
+        self.assertTrue(len(stderr) > 0)
+        self.assertTrue(user in stderr[0])
+
+    def test_run_command_shell(self):
+        output = self.client.run_command(self.fake_cmd, shell="bash -c")
+        self.client.join(output)
+        stdout = list(output[self.host].stdout)
+        self.assertEqual(stdout, [self.fake_resp])
+
+    def test_run_command_no_shell(self):
+        output = self.client.run_command('id', use_shell=False)
+        self.client.join(output)
+        stdout = list(output[self.host].stdout)
+        self.assertTrue(len(stdout) > 0)
+        self.assertTrue(output[self.host].exit_code == 0)
+
+    def test_run_command_environment(self):
+        env = {'ENV_VARIABLE': 'env value'}
+        output = self.client.run_command('echo ${ENV_VARIABLE}',
+                                         environment=env)
+        self.client.join(output)
+        stdout = list(output[self.host].stdout)
+        expected = [env.values()[0]]
+        self.assertEqual(stdout, expected)
+
 if __name__ == '__main__':
     unittest.main()