Merge "Fix incorrect exit status checking in sshutils"

rally/sshutils.py

@@ -93,35 +93,34 @@ class SSH(object):
         self._get_ssh_connection()
         cmd = ' '.join(cmd)
         transport = self.client.get_transport()
-        channel = transport.open_session()
-        channel.fileno()
-        channel.exec_command(cmd)
-        channel.shutdown_write()
-        poll = select.poll()
-        poll.register(channel, select.POLLIN)
+        session = transport.open_session()
+        session.exec_command(cmd)
         start_time = time.time()
+
         while True:
-            ready = poll.poll(16)
-            if not any(ready):
-                if not self._is_timed_out(start_time):
-                    continue
-                raise exceptions.TimeoutException('SSH Timeout')
-            if not ready[0]:
-                continue
-            out_chunk = err_chunk = None
-            if channel.recv_ready():
-                out_chunk = channel.recv(4096)
+            errors = select.select([session], [], [], 4)[2]
+
+            if session.recv_ready():
+                data = session.recv(4096)
+                LOG.debug(data)
                 if get_stdout:
-                    yield (1, out_chunk)
-                LOG.debug("stdout: %s" % out_chunk)
-            if channel.recv_stderr_ready():
-                err_chunk = channel.recv_stderr(4096)
+                    yield (1, data)
+                continue
+
+            if session.recv_stderr_ready():
+                data = session.recv_stderr(4096)
+                LOG.debug(data)
                 if get_stderr:
-                    yield (2, err_chunk)
-                LOG.debug("stderr: %s" % err_chunk)
-            if channel.closed and not err_chunk and not out_chunk:
+                    yield (2, data)
+                continue
+
+            if errors or session.exit_status_ready():
                 break
-        exit_status = channel.recv_exit_status()
+
+            if self._is_timed_out(start_time):
+                raise exceptions.TimeoutException('SSH Timeout')
+
+        exit_status = session.recv_exit_status()
         if 0 != exit_status:
             raise exceptions.SSHError(
                 'SSHExecCommandFailed with exit_status %s'

tests/test_sshutils.py

@@ -27,49 +27,49 @@ class SSHTestCase(test.TestCase):
         super(SSHTestCase, self).setUp()
         self.ssh = sshutils.SSH('example.net', 'root')
         self.channel = mock.Mock()
-        self.channel.fileno.return_value = 15
-        self.channel.recv.return_value = ''
-        self.channel.recv_stderr.return_value = ''
+        self.channel.recv.return_value = 'ok'
+        self.channel.recv_stderr.return_value = 'error'
         self.channel.recv_exit_status.return_value = 0
         self.transport = mock.Mock()
         self.transport.open_session = mock.MagicMock(return_value=self.channel)
-        self.poll = mock.Mock()
-        self.poll.poll.return_value = [(self.channel, 1)]
         self.policy = mock.Mock()
         self.client = mock.Mock()
         self.client.get_transport = mock.MagicMock(return_value=self.transport)
 
+        self.channel.exit_status_ready.return_value = True
+        self.channel.recv_ready.side_effect = [True, False, False]
+        self.channel.recv_stderr_ready.side_effect = [True, False, False]
+
     @mock.patch('rally.sshutils.paramiko')
     @mock.patch('rally.sshutils.select')
     def test_generator(self, st, pk):
         pk.SSHClient.return_value = self.client
-        st.poll.return_value = self.poll
-
-        self.channel.recv.side_effect = ['ok', '']
-        self.channel.recv_stderr.side_effect = ['error', '']
+        st.select.return_value = ([], [], [])
 
         chunks = list(self.ssh.execute_generator('ps ax'))
-        self.assertEqual([(1, 'ok'), (2, 'error'), (1, ''), (2, '')], chunks)
+        self.assertEqual([(1, 'ok'), (2, 'error')], chunks)
 
     @mock.patch('rally.sshutils.paramiko')
     @mock.patch('rally.sshutils.select')
     def test_execute(self, st, pk):
         pk.SSHClient.return_value = self.client
-        st.poll.return_value = self.poll
+        st.select.return_value = ([], [], [])
         stdout, stderr = self.ssh.execute('uname')
 
         self.assertEqual('', stdout)
         self.assertEqual('', stderr)
-        expected = [mock.call.fileno(),
-                    mock.call.exec_command('uname'),
-                    mock.call.shutdown_write(),
+        expected = [mock.call.exec_command('uname'),
                     mock.call.recv_ready(),
                     mock.call.recv(4096),
+                    mock.call.recv_ready(),
                     mock.call.recv_stderr_ready(),
                     mock.call.recv_stderr(4096),
+                    mock.call.recv_ready(),
+                    mock.call.recv_stderr_ready(),
+                    mock.call.exit_status_ready(),
                     mock.call.recv_exit_status()]
 
-        self.assertEqual(self.channel.mock_calls, expected)
+        self.assertEqual(expected, self.channel.mock_calls)
 
     @mock.patch('rally.sshutils.paramiko')
     def test_upload_file(self, pk):