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 <ross.b.brattain@intel.com>

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
@@ -275,6 +277,23 @@ class SSHRunTestCase(unittest.TestCase):
         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
         mock_select.select.return_value = ([], [], [True])
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