From a8db6ed817a39632bed19e0ece64f1862231db16 Mon Sep 17 00:00:00 2001 From: Ross Brattain Date: Sun, 13 Nov 2016 18:08:57 -0800 Subject: ssh.py: add flag to keep stdin open For some VNFs we may want to send periodic commands, for example to print statistics, but otherwise not write anything for long periods of time. Currently when we can no longer read from stdin we close it. A workaround is to constantly spam stdin with newlines to keep forcing stdin open. We don't want to have to do this, so add an enable flag to keep stdin open. If the caller wants to close stdin at some point it can. Change-Id: I9496022295dfd19804572e484fe4f170ca7d4ac3 Signed-off-by: Ross Brattain --- tests/unit/test_ssh.py | 19 +++++++++++++++++++ yardstick/ssh.py | 28 ++++++++++++++++++---------- 2 files changed, 37 insertions(+), 10 deletions(-) diff --git a/tests/unit/test_ssh.py b/tests/unit/test_ssh.py index a27052462..1e021a051 100644 --- a/tests/unit/test_ssh.py +++ b/tests/unit/test_ssh.py @@ -18,6 +18,8 @@ import os import unittest +from cStringIO import StringIO + import mock from yardstick import ssh @@ -274,6 +276,23 @@ class SSHRunTestCase(unittest.TestCase): send_calls = [call("line1"), call("line2"), call("e2")] self.assertEqual(send_calls, self.fake_session.send.mock_calls) + @mock.patch("yardstick.ssh.select") + def test_run_stdin_keep_open(self, mock_select): + """Test run method with stdin. + + Third send call was called with "e2" because only 3 bytes was sent + by second call. So remainig 2 bytes of "line2" was sent by third call. + """ + mock_select.select.return_value = ([], [], []) + self.fake_session.exit_status_ready.side_effect = [0, 0, 0, True] + self.fake_session.send_ready.return_value = True + self.fake_session.send.side_effect = len + fake_stdin = StringIO("line1\nline2\n") + self.test_client.run("cmd", stdin=fake_stdin, keep_stdin_open=True) + call = mock.call + send_calls = [call("line1\nline2\n")] + self.assertEqual(send_calls, self.fake_session.send.mock_calls) + @mock.patch("yardstick.ssh.select") def test_run_select_error(self, mock_select): self.fake_session.exit_status_ready.return_value = False diff --git a/yardstick/ssh.py b/yardstick/ssh.py index 8b71fe606..e0e2f83ee 100644 --- a/yardstick/ssh.py +++ b/yardstick/ssh.py @@ -140,10 +140,12 @@ class SSH(object): self._client = False def run(self, cmd, stdin=None, stdout=None, stderr=None, - raise_on_error=True, timeout=3600): + raise_on_error=True, timeout=3600, + keep_stdin_open=False): """Execute specified command on the server. :param cmd: Command to be executed. + :type cmd: str :param stdin: Open file or string to pass to stdin. :param stdout: Open file to connect to stdout. :param stderr: Open file to connect to stderr. @@ -151,6 +153,8 @@ class SSH(object): then exception will be raized if non-zero code. :param timeout: Timeout in seconds for command execution. Default 1 hour. No timeout if set to 0. + :param keep_stdin_open: don't close stdin on empty reads + :type keep_stdin_open: bool """ client = self._get_client() @@ -160,10 +164,12 @@ class SSH(object): return self._run(client, cmd, stdin=stdin, stdout=stdout, stderr=stderr, raise_on_error=raise_on_error, - timeout=timeout) + timeout=timeout, + keep_stdin_open=keep_stdin_open) def _run(self, client, cmd, stdin=None, stdout=None, stderr=None, - raise_on_error=True, timeout=3600): + raise_on_error=True, timeout=3600, + keep_stdin_open=False): transport = client.get_transport() session = transport.open_session() @@ -203,13 +209,15 @@ class SSH(object): if not data_to_send: data_to_send = stdin.read(4096) if not data_to_send: - stdin.close() - session.shutdown_write() - writes = [] - continue - sent_bytes = session.send(data_to_send) - # LOG.debug("sent: %s" % data_to_send[:sent_bytes]) - data_to_send = data_to_send[sent_bytes:] + # we may need to keep stdin open + if not keep_stdin_open: + stdin.close() + session.shutdown_write() + writes = [] + if data_to_send: + sent_bytes = session.send(data_to_send) + # LOG.debug("sent: %s" % data_to_send[:sent_bytes]) + data_to_send = data_to_send[sent_bytes:] if session.exit_status_ready(): break -- cgit 1.2.3-korg