aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--tests/unit/test_ssh.py60
-rw-r--r--yardstick/ssh.py34
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)