Merge "import new _put_file_shell method from upstream rally"

2 files changed, 92 insertions, 2 deletions

diff --git a/tests/unit/test_ssh.py b/tests/unit/test_ssh.py
index 1e021a051..88638a0a8 100644
--- a/tests/unit/test_ssh.py
+++ b/tests/unit/test_ssh.py
@@ -17,6 +17,7 @@
 # rally/tests/unit/common/test_sshutils.py
 
 import os
+import socket
 import unittest
 from cStringIO import StringIO
 
@@ -164,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):
@@ -307,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()
diff --git a/yardstick/ssh.py b/yardstick/ssh.py
index 8485dccd0..b9d9262cf 100644
--- a/yardstick/ssh.py
+++ b/yardstick/ssh.py
@@ -290,3 +290,37 @@ class SSH(object):
     def send_command(self, command):
         client = self._get_client()
         client.exec_command(command, get_pty=True)
+
+    def _put_file_sftp(self, localpath, remotepath, mode=None):
+        client = self._get_client()
+
+        with client.open_sftp() as sftp:
+            sftp.put(localpath, remotepath)
+            if mode is None:
+                mode = 0o777 & os.stat(localpath).st_mode
+            sftp.chmod(remotepath, mode)
+
+    def _put_file_shell(self, localpath, remotepath, mode=None):
+        # quote to stop wordpslit
+        cmd = ['cat > "%s"' % remotepath]
+        if mode is not None:
+            # use -- so no options
+            cmd.append('chmod -- 0%o "%s"' % (mode, remotepath))
+
+        with open(localpath, "rb") as localfile:
+            # only chmod on successful cat
+            cmd = "&& ".join(cmd)
+            self.run(cmd, stdin=localfile)
+
+    def put_file(self, localpath, remotepath, mode=None):
+        """Copy specified local file to the server.
+
+        :param localpath:   Local filename.
+        :param remotepath:  Remote filename.
+        :param mode:        Permissions to set after upload
+        """
+        import socket
+        try:
+            self._put_file_sftp(localpath, remotepath, mode=mode)
+        except (paramiko.SSHException, socket.error):
+            self._put_file_shell(localpath, remotepath, mode=mode)