Added timeout ability to fake server. Added working test for timeout argument of ssh clients.

Author: Panos Kittenis

fake_server/fake_server.py

@@ -26,9 +26,8 @@
 """
 
 import gevent
-from gevent import monkey
-monkey.patch_all()
 import os
+import socket
 from gevent import socket
 from gevent.event import Event
 import sys
@@ -37,14 +36,17 @@
 import paramiko
 import time
 from stub_sftp import StubSFTPServer
+from gevent import monkey
+monkey.patch_all()
+
 
 logger = logging.getLogger("fake_server")
 paramiko_logger = logging.getLogger('paramiko.transport')
 
 host_key = paramiko.RSAKey(filename = os.path.sep.join([os.path.dirname(__file__), 'rsa.key']))
 
 class Server (paramiko.ServerInterface):
-    def __init__(self, cmd_req_response = {}, fail_auth = False):
+    def __init__(self, cmd_req_response = {}, fail_auth=False):
         self.event = Event()
         self.cmd_req_response = cmd_req_response
         self.fail_auth = fail_auth
@@ -113,7 +115,8 @@ def make_socket(listen_ip):
         return
     return sock
 
-def listen(cmd_req_response, sock, fail_auth = False):
+def listen(cmd_req_response, sock, fail_auth=False,
+           timeout=None):
     """Run a fake ssh server and given a cmd_to_run, send given \
     response to client connection. Returns (server, socket) tuple \
     where server is a joinable server thread and socket is listening \
@@ -129,16 +132,17 @@ def listen(cmd_req_response, sock, fail_auth = False):
         logger.error('*** Listen failed: %s' % (str(e),))
         traceback.print_exc()
         return
-    handle_ssh_connection(cmd_req_response, sock, fail_auth=fail_auth)
+    handle_ssh_connection(cmd_req_response, sock, fail_auth=fail_auth,
+                          timeout=timeout)
 
-def _handle_ssh_connection(cmd_req_response, transport, fail_auth = False):
+def _handle_ssh_connection(cmd_req_response, transport, fail_auth=False):
     try:
         transport.load_server_moduli()
     except:
         return
     transport.add_server_key(host_key)
     transport.set_subsystem_handler('sftp', paramiko.SFTPServer, StubSFTPServer)
-    server = Server(cmd_req_response = cmd_req_response, fail_auth = fail_auth)
+    server = Server(cmd_req_response = cmd_req_response, fail_auth=fail_auth)
     try:
         transport.start_server(server=server)
     except paramiko.SSHException, e:
@@ -158,9 +162,15 @@ def _handle_ssh_connection(cmd_req_response, transport, fail_auth = False):
         time.sleep(.5)
     channel.close()
 
-def handle_ssh_connection(cmd_req_response, sock, fail_auth = False):
+def handle_ssh_connection(cmd_req_response, sock,
+                          fail_auth=False,
+                          timeout=None):
     conn, addr = sock.accept()
     logger.info('Got connection..')
+    if timeout:
+        logger.debug("SSH server sleeping for %s then raising socket.timeout",
+                    timeout)
+        gevent.Timeout(timeout).start()
     try:
         transport = paramiko.Transport(conn)
         _handle_ssh_connection(cmd_req_response, transport, fail_auth=fail_auth)
@@ -173,8 +183,10 @@ def handle_ssh_connection(cmd_req_response, sock, fail_auth = False):
             pass
         return
 
-def start_server(cmd_req_response, sock, fail_auth=False):
-    return gevent.spawn(listen, cmd_req_response, sock, fail_auth=fail_auth)
+def start_server(cmd_req_response, sock, fail_auth=False,
+                 timeout=None):
+    return gevent.spawn(listen, cmd_req_response, sock, fail_auth=fail_auth,
+                        timeout=timeout)
 
 if __name__ == "__main__":
     logging.basicConfig()

pssh.py

@@ -66,6 +66,11 @@ class ProxyCommandException(Exception):
     pass
 
 
+class SSHException(Exception):
+    """Raised on SSHException error - error authenticating with SSH server"""
+    pass
+
+
 class SSHClient(object):
     """Wrapper class over paramiko.SSHClient with sane defaults
     Honours ~/.ssh/config and /etc/ssh/ssh_config entries for host username \
@@ -178,6 +183,10 @@ def _connect(self, retries=1):
         except paramiko.ProxyCommandFailure, e:
             logger.error("Error executing ProxyCommand - %s", e.message,)
             raise ProxyCommandException(e.message)
+        # SSHException is more general so should below other types
+        # of SSH failure
+        except paramiko.SSHException, e:
+            raise SSHException(e)
 
     def exec_command(self, command, sudo=False, user=None, **kwargs):
         """Wrapper to :mod:`paramiko.SSHClient.exec_command`
@@ -316,11 +325,11 @@ def __init__(self, hosts,
         :param pool_size: (Optional) Greenlet pool size. Controls on how many\
         hosts to execute tasks in parallel. Defaults to 10
         :type pool_size: int
-
+        
         **Example**
 
         >>> from pssh import ParallelSSHClient, AuthenticationException,\
-                UnknownHostException, ConnectionErrorException
+        		UnknownHostException, ConnectionErrorException
         >>> client = ParallelSSHClient(['myhost1', 'myhost2'])
         >>> try:
         >>> ... cmds = client.exec_command('ls -ltrh /tmp/aasdfasdf', sudo = True)
@@ -335,7 +344,7 @@ def __init__(self, hosts,
         **Example with returned stdout and stderr buffers**
 
         >>> from pssh import ParallelSSHClient, AuthenticationException,\
-                UnknownHostException, ConnectionErrorException
+        		UnknownHostException, ConnectionErrorException
         >>> client = ParallelSSHClient(['myhost1', 'myhost2'])
         >>> try:
         >>> ... cmds = client.exec_command('ls -ltrh /tmp/aasdfasdf', sudo = True)
@@ -344,10 +353,10 @@ def __init__(self, hosts,
         >>> output = [client.get_stdout(cmd, return_buffers=True) for cmd in cmds]
         >>> print output
         [{'myhost1': {'exit_code': 2,
-                  'stdout' : <generator object <genexpr>,
+                      'stdout' : <generator object <genexpr>,
                       'stderr' : <generator object <genexpr>,}},
          {'myhost2': {'exit_code': 2,
-                  'stdout' : <generator object <genexpr>,
+                      'stdout' : <generator object <genexpr>,
                       'stderr' : <generator object <genexpr>,}},
                       ]
         >>> for host_stdout in output:
@@ -361,25 +370,25 @@ def __init__(self, hosts,
         >>> import paramiko
         >>> client_key = paramiko.RSAKey.from_private_key_file('user.key')
         >>> client = ParallelSSHClient(['myhost1', 'myhost2'], pkey=client_key)
-
+        
         .. note ::
 
           **Connection persistence**
-
+          
           Connections to hosts will remain established for the duration of the
           object's life. To close them, just `del` or reuse the object reference.
-
+          
           >>> client = ParallelSSHClient(['localhost'])
           >>> cmds = client.exec_command('ls -ltrh /tmp/aasdfasdf')
           >>> cmds[0].join()
-
+          
           :netstat: ``tcp        0      0 127.0.0.1:53054         127.0.0.1:22            ESTABLISHED``
-
+          
           Connection remains active after commands have finished executing. Any \
           additional commands will use the same connection.
-
+          
           >>> del client
-
+          
           Connection is terminated.
         """
         self.pool = gevent.pool.Pool(size=pool_size)
@@ -408,19 +417,19 @@ def exec_command(self, *args, **kwargs):
         **Example**:
 
         >>> cmds = client.exec_command('ls -ltrh')
-
+        
         Wait for completion, no stdout:
-
+        
         >>> for cmd in cmds:
         >>>     cmd.join()
-
+        
         Alternatively/in addition print stdout for each command:
-
+        
         >>> print [get_stdout(cmd) for cmd in cmds]
 
         Retrieving stdout implies join, meaning get_stdout will wait
         for completion of all commands before returning output.
-
+        
         You may call get_stdout on already completed greenlets to re-get
         their output as many times as you want."""
         return [self.pool.spawn(self._exec_command, host, *args, **kwargs)
@@ -439,7 +448,7 @@ def _exec_command(self, host, *args, **kwargs):
 
     def get_stdout(self, greenlet, return_buffers=False):
         """Get/print stdout from greenlet and return exit code for host
-
+        
         :mod:`pssh.get_stdout` will close the open SSH channel but this does
         **not** close the established connection to the remote host, only the
         authenticated SSH channel within it. This is standard practise
@@ -463,7 +472,7 @@ def get_stdout(self, greenlet, return_buffers=False):
         for example ``{'myhost1': {'exit_code': 0}}``
         :rtype: With ``return_buffers=True``: ``{'myhost1': {'exit_code': 0,
                                                              'channel' : None or SSH channel of command if command is still executing,
-                                     'stdout' : <iterable>,
+                                                             'stdout' : <iterable>,
                                                              'stderr' : <iterable>,}}``
         """
         gevent.sleep(.2)

tests/test_pssh_client.py

@@ -111,6 +111,27 @@ def test_pssh_client_auth_failure(self):
         del client
         server.join()
 
+    def test_pssh_client_timeout(self):
+        server = start_server({ self.fake_cmd : self.fake_resp },
+                              self.listen_socket,
+                              timeout=0.2)
+        client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
+                                   pkey=self.user_key,
+                                   timeout=0.1)
+        cmd = client.exec_command(self.fake_cmd)[0]
+        # Handle exception
+        try:
+            gevent.sleep(0.5)
+            cmd.get()
+            if not server.exception:
+                raise Exception("Expected EOFError from socket timeout, got none")
+            raise server.exception
+        except gevent.Timeout:
+            pass
+        del client
+        server.join()
+
+
     def test_pssh_client_exec_command_password(self):
         """Test password authentication. Fake server accepts any password
         even empty string"""

tests/test_ssh_client.py

@@ -20,11 +20,14 @@
 
 """Unittests for :mod:`pssh.SSHClient` class"""
 
+import gevent
+import socket
+import time
 import unittest
 from pssh import SSHClient, ParallelSSHClient, UnknownHostException, AuthenticationException,\
-    logger, ConnectionErrorException, UnknownHostException
+     logger, ConnectionErrorException, UnknownHostException, SSHException
 from fake_server.fake_server import start_server, make_socket, logger as server_logger, \
-    paramiko_logger
+     paramiko_logger
 from fake_server.fake_agent import FakeAgent
 import paramiko
 import os
@@ -126,20 +129,13 @@ def test_ssh_client_retries(self):
                          msg="Got unexpected number of retries %s - expected %s"
                          % (num_tries, expected_num_tries,))
 
-    def test_ssh_client_conn_failure(self):
+    def test_ssh_client_unknown_host_failure(self):
         """Test connection error failure case - ConnectionErrorException"""
         host = ''.join([random.choice(string.ascii_letters) for n in xrange(8)])
         self.assertRaises(UnknownHostException,
                           SSHClient, host, port=self.listen_port,
                           pkey=self.user_key, num_retries=0)
 
-    def test_ssh_client_timeout(self):
-        """Test connection timeout error"""
-        with self.assertRaises(ConnectionErrorException) as cm:
-            SSHClient('127.0.0.1', port=self.listen_port,
-                      pkey=self.user_key, num_retries=0, timeout=1)
-
-        self.assertEqual(cm.exception.args[1], 'timed out')
 
 if __name__ == '__main__':
     unittest.main()