1 files changed, 77 insertions, 2 deletions

diff --git a/tests/unit/test_ssh.py b/tests/unit/test_ssh.py
index a27052462..8b828ed7c 100644
--- a/tests/unit/test_ssh.py
+++ b/tests/unit/test_ssh.py
@@ -17,7 +17,10 @@
 # rally/tests/unit/common/test_sshutils.py
 
 import os
+import socket
 import unittest
+from cStringIO import StringIO
+
 import mock
 
 from yardstick import ssh
@@ -162,10 +165,10 @@ class SSHTestCase(unittest.TestCase):
     def test_send_command(self, mock_paramiko):
         paramiko_sshclient = self.test_client._get_client()
         with mock.patch.object(paramiko_sshclient, "exec_command") \
-            as mock_paramiko_exec_command:
+                as mock_paramiko_exec_command:
             self.test_client.send_command('cmd')
         mock_paramiko_exec_command.assert_called_once_with('cmd',
-                                                            get_pty=True)
+                                                           get_pty=True)
 
 
 class SSHRunTestCase(unittest.TestCase):
@@ -275,6 +278,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])
@@ -288,6 +308,61 @@ class SSHRunTestCase(unittest.TestCase):
         self.fake_session.exit_status_ready.return_value = False
         self.assertRaises(ssh.SSHTimeout, self.test_client.run, "cmd")
 
+    @mock.patch("yardstick.ssh.open", create=True)
+    def test__put_file_shell(self, mock_open):
+        self.test_client.run = mock.Mock()
+        self.test_client._put_file_shell("localfile", "remotefile", 0o42)
+
+        self.test_client.run.assert_called_once_with(
+            'cat > remotefile && chmod -- 042 remotefile',
+            stdin=mock_open.return_value.__enter__.return_value)
+
+    @mock.patch("yardstick.ssh.os.stat")
+    def test__put_file_sftp(self, mock_stat):
+        sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
+        sftp.__enter__.return_value = sftp
+
+        mock_stat.return_value = os.stat_result([0o753] + [0] * 9)
+
+        self.test_client._put_file_sftp("localfile", "remotefile")
+
+        sftp.put.assert_called_once_with("localfile", "remotefile")
+        mock_stat.assert_called_once_with("localfile")
+        sftp.chmod.assert_called_once_with("remotefile", 0o753)
+        sftp.__exit__.assert_called_once_with(None, None, None)
+
+    def test__put_file_sftp_mode(self):
+        sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
+        sftp.__enter__.return_value = sftp
+
+        self.test_client._put_file_sftp("localfile", "remotefile", mode=0o753)
+
+        sftp.put.assert_called_once_with("localfile", "remotefile")
+        sftp.chmod.assert_called_once_with("remotefile", 0o753)
+        sftp.__exit__.assert_called_once_with(None, None, None)
+
+    def test_put_file_SSHException(self):
+        exc = ssh.paramiko.SSHException
+        self.test_client._put_file_sftp = mock.Mock(side_effect=exc())
+        self.test_client._put_file_shell = mock.Mock()
+
+        self.test_client.put_file("foo", "bar", 42)
+        self.test_client._put_file_sftp.assert_called_once_with("foo", "bar",
+                                                                mode=42)
+        self.test_client._put_file_shell.assert_called_once_with("foo", "bar",
+                                                                 mode=42)
+
+    def test_put_file_socket_error(self):
+        exc = socket.error
+        self.test_client._put_file_sftp = mock.Mock(side_effect=exc())
+        self.test_client._put_file_shell = mock.Mock()
+
+        self.test_client.put_file("foo", "bar", 42)
+        self.test_client._put_file_sftp.assert_called_once_with("foo", "bar",
+                                                                mode=42)
+        self.test_client._put_file_shell.assert_called_once_with("foo", "bar",
+                                                                 mode=42)
+
 
 def main():
     unittest.main()