Made fake server port an optional parameter. Added test for expected output behaviour when pool size is less than number of hosts - resolves #25

Author: Dan

fake_server/fake_server.py

@@ -119,12 +119,12 @@ def _long_running_response(self, channel, responder):
         channel.send_exit_status(0)
         channel.close()
 
-def make_socket(listen_ip):
+def make_socket(listen_ip, port=0):
     """Make socket on given address and available port chosen by OS"""
     try:
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        sock.bind((listen_ip, 0))
+        sock.bind((listen_ip, port))
     except Exception, e:
         logger.error('Failed to bind to address - %s' % (str(e),))
         traceback.print_exc()

pssh_local.py

@@ -44,14 +44,13 @@ def test():
 
 def test_parallel():
     """Perform ls and copy file with ParallelSSHClient on localhost"""
-    client = ParallelSSHClient(['localhost'])
-    cmds = client.exec_command('ls -ltrh')
-    output = [client.get_stdout(cmd, return_buffers=True) for cmd in cmds]
+    client = ParallelSSHClient(['localhost', 'localhost'])
+    output = client.run_command('ls -ltrh')
     print output
-    cmds = client.copy_file('../test', 'test_dir/test')
-    client.pool.join()
+    # cmds = client.copy_file('../test', 'test_dir/test')
+    # client.pool.join()
 
 if __name__ == "__main__":
     _setup_logger(logger)
-    test()
+    # test()
     test_parallel()

tests/test_pssh_client.py

@@ -231,7 +231,7 @@ def test_pssh_client_exec_command_password(self):
                          msg="Got unexpected command output - %s" % (output,))
         del client
         server.join()
-                
+
     def test_pssh_client_long_running_command(self):
         expected_lines = 5
         server = start_server({ self.long_running_cmd :
@@ -314,8 +314,32 @@ def test_pssh_pool_size(self):
     def test_pssh_hosts_more_than_pool_size(self):
         """Test we can successfully run on more hosts than our pool size and
         get logs for all hosts"""
-        raise NotImplementedError
-        
+        # Make a second server on the same port as the first one
+        server2_socket = make_socket('127.0.0.2', port=self.listen_port)
+        server2_port = server2_socket.getsockname()[1]
+        server1 = start_server({ self.fake_cmd : self.fake_resp },
+                               self.listen_socket)
+        server2 = start_server({ self.fake_cmd : self.fake_resp },
+                               server2_socket)
+        hosts = ['127.0.0.1', '127.0.0.2']
+        client = ParallelSSHClient(hosts,
+                                   port=self.listen_port,
+                                   pkey=self.user_key,
+                                   pool_size=1,
+                                   )
+        output = client.run_command(self.fake_cmd)
+        stdout = [list(output[k]['stdout']) for k in output]
+        expected_stdout = [[self.fake_resp], [self.fake_resp]]
+        self.assertEqual(len(hosts), len(output),
+                         msg="Did not get output from all hosts. Got output for \
+%s/%s hosts" % (len(output), len(hosts),))
+        self.assertEqual(expected_stdout, stdout,
+                         msg="Did not get expected output from all hosts. \
+                         Got %s - expected %s" % (stdout, expected_stdout,))
+        del client
+        server1.kill()
+        server2.kill()
+
     def test_ssh_proxy(self):
         """Test connecting to remote destination via SSH proxy
         client -> proxy -> destination