Updated tests to be python 2.6 compatible

ParallelSSH · Jan 9, 2015 · 4d2035f · 4d2035f
1 parent ecf56fc
commit 4d2035f
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 26 deletions.
diff --git a/tests/test_pssh_client.py b/tests/test_pssh_client.py
@@ -252,16 +252,19 @@ def test_pssh_client_long_running_command(self):
     def test_pssh_client_retries(self):
         """Test connection error retries"""
         expected_num_tries = 2
-        with self.assertRaises(ConnectionErrorException) as cm:
-            client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
-                                       pkey=self.user_key,
-                                       num_retries=expected_num_tries)
-            cmd = client.exec_command('blah')[0]
-            cmd.get()
-        num_tries = cm.exception.args[-1:][0]
-        self.assertEqual(expected_num_tries, num_tries,
-                         msg="Got unexpected number of retries %s - expected %s"
-                         % (num_tries, expected_num_tries,))
+        client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
+                                   pkey=self.user_key,
+                                   num_retries=expected_num_tries)
+        self.assertRaises(ConnectionErrorException, client.run_command, 'blah')
+        try:
+            client.run_command('blah')
+        except ConnectionErrorException, ex:
+            num_tries = ex.args[-1:][0]
+            self.assertEqual(expected_num_tries, num_tries,
+                             msg="Got unexpected number of retries %s - expected %s"
+                             % (num_tries, expected_num_tries,))
+        else:
+            raise Exception('No ConnectionErrorException')
 
     def test_pssh_copy_file(self):
         """Test parallel copy file"""

diff --git a/tests/test_ssh_client.py b/tests/test_ssh_client.py
@@ -111,23 +111,12 @@ def test_ssh_client_conn_failure(self):
                           pkey=self.user_key, num_retries=0)
 
     def test_ssh_client_retries(self):
-        """Test connection error retries"""
-        expected_num_tries = 2
-        with self.assertRaises(ConnectionErrorException) as cm:
-            SSHClient('127.0.0.1', port=self.listen_port,
-                      pkey=self.user_key, num_retries=expected_num_tries)
-        num_tries = cm.exception.args[-1:][0]
-        self.assertEqual(expected_num_tries, num_tries,
-                         msg="Got unexpected number of retries %s - expected %s"
-                         % (num_tries, expected_num_tries,))
+        """Test connection error exceptions"""
+        self.assertRaises(ConnectionErrorException, SSHClient, '127.0.0.1', port=self.listen_port,
+                          pkey=self.user_key, num_retries=1)
         host = ''.join([random.choice(string.ascii_letters) for n in xrange(8)])
-        with self.assertRaises(UnknownHostException) as cm:
-            SSHClient(host, port=self.listen_port,
-                      pkey=self.user_key, num_retries=expected_num_tries)
-        num_tries = cm.exception.args[-1:][0]
-        self.assertEqual(expected_num_tries, num_tries,
-                         msg="Got unexpected number of retries %s - expected %s"
-                         % (num_tries, expected_num_tries,))
+        self.assertRaises(UnknownHostException, SSHClient, host, port=self.listen_port,
+                          pkey=self.user_key, num_retries=1)
 
     def test_ssh_client_unknown_host_failure(self):
         """Test connection error failure case - ConnectionErrorException"""