Added join convenience function to parallel client and unittest

pssh/pssh_client.py

@@ -387,7 +387,15 @@ def _update_host_output(self, output, host, exit_code, channel, stdout, stderr,
                              'stderr' : stderr,
                              'cmd' : cmd,
                              'exception' : exception,})
-
+
+    def join(self, output):
+        """Block until all remote commands in output have finished
+        and retrieve exit codes"""
+        for host in output:
+            for line in output[host]['stdout']:
+                pass
+        self.get_exit_codes(output)
+
     def get_exit_codes(self, output):
         """Get exit code for all hosts in output *if available*.
         Output parameter is modified in-place.

tests/test_pssh_client.py

@@ -293,8 +293,7 @@ def test_pssh_client_long_running_command_exit_codes(self):
         # Embedded server is also asynchronous and in the same thread
         # as our client so need to sleep for duration of server connection
         gevent.sleep(expected_lines)
-        client.pool.join()
-        client.get_exit_codes(output)
+        client.join(output)
         self.assertTrue(output[self.host]['exit_code'] == 0,
                         msg="Got non-zero exit code %s" % (
                             output[self.host]['exit_code'],))