Added retries on connection failures to ssh client. Changed tests to wait for server to come up before client trying to connect

Author: pkittenis

pssh.py

@@ -17,6 +17,7 @@
 handler.setFormatter(host_log_format)
 host_logger.addHandler(handler)
 host_logger.setLevel(logging.INFO)
+NUM_RETRIES = 3
 
 logger = logging.getLogger(__name__)
 
@@ -90,24 +91,29 @@ def __init__(self, host,
         self.host = resolved_address
         self._connect()
 
-    def _connect(self):
+    def _connect(self, retries=1):
         """Connect to host, throw UnknownHost exception on DNS errors"""
         try:
             self.client.connect(self.host, username=self.user,
                                 password=self.password, port=self.port)
         except socket.gaierror, e:
             logger.error("Could not resolve host '%s'", self.host,)
+            while retries < NUM_RETRIES:
+                gevent.sleep(5)
+                return self._connect(retries=retries+1)
             raise UnknownHostException("%s - %s" % (str(e.args[1]),
                                                     self.host,))
         except socket.error, e:
             logger.error("Error connecting to host '%s:%s'" % (self.host,
                                                                self.port,))
+            while retries < NUM_RETRIES:
+                gevent.sleep(5)
+                return self._connect(retries=retries+1)
             raise ConnectionErrorException("%s for host '%s:%s'" % (str(e.args[1]),
                                                                     self.host,
                                                                     self.port,))
         except paramiko.AuthenticationException, e:
             raise AuthenticationException(e)
-        gevent.sleep(0)
 
     def exec_command(self, command, sudo=False, **kwargs):
         """Wrapper to :mod:`paramiko.SSHClient.exec_command`
@@ -147,7 +153,6 @@ def _make_sftp(self):
         """Make SFTP client from open transport"""
         transport = self.client.get_transport()
         channel = transport.open_session()
-        gevent.sleep(0)
         return paramiko.SFTPClient.from_transport(transport)
 
     def mkdir(self, sftp, directory):
@@ -164,7 +169,6 @@ def mkdir(self, sftp, directory):
         except IOError, error:
             logger.error("Error occured creating directory on %s - %s",
                          self.host, error)
-        gevent.sleep(0)
 
     def copy_file(self, local_file, remote_file):
         """Copy local file to host via SFTP/SCP
@@ -195,7 +199,6 @@ def copy_file(self, local_file, remote_file):
         else:
             logger.info("Copied local file %s to remote destination %s:%s",
                         local_file, self.host, remote_file)
-        gevent.sleep(0)
 
 class ParallelSSHClient(object):
     """

tests/test_pssh_client.py

@@ -10,6 +10,7 @@
 import logging
 import gevent
 import threading
+import time
 
 _setup_logger(server_logger)
 
@@ -34,7 +35,28 @@ def _run_client_exec_command(self):
         gevent.sleep(0)
         del client
 
-    def _run_client_fail_auth(self):
+    # def _run_client_fail_auth(self):
+    #     try:
+    #     except Exception:
+    #         raise SystemExit
+
+    def test_pssh_client_exec_command(self):
+        server = listen({ self.fake_cmd : self.fake_resp }, self.listener)
+        time.sleep(5)
+        client = ParallelSSHClient(['localhost'], port=self.listen_port)
+        cmd = client.exec_command(self.fake_cmd)[0]
+        output = client.get_stdout(cmd)
+        expected = {'localhost' : {'exit_code' : 0}}
+        self.assertEqual(expected, output,
+                         msg = "Got unexpected command output - %s" % (output,))
+        gevent.sleep(0)
+        del client
+        server.join()
+
+    def test_pssh_client_auth_failure(self):
+        server = listen({ self.fake_cmd : self.fake_resp },
+                        self.listener, fail_auth=True)
+        time.sleep(5)
         client = ParallelSSHClient(['localhost'], port=self.listen_port)
         cmd = client.exec_command(self.fake_cmd)[0]
         # Handle exception
@@ -44,18 +66,4 @@ def _run_client_fail_auth(self):
         except AuthenticationException:
             pass
         del client
-
-    def test_pssh_client_exec_command(self):
-        server = listen({ self.fake_cmd : self.fake_resp }, self.listener)
-        client = threading.Thread(target=self._run_client_exec_command)
-        client.start()
-        client.join()
-        server.join()
-
-    def test_pssh_client_auth_failure(self):
-        server = listen({ self.fake_cmd : self.fake_resp },
-                        self.listener, fail_auth=True)
-        client = threading.Thread(target=self._run_client_fail_auth)
-        client.start()
-        client.join()
         server.join()