From 6a49981bc5d8f9f6ca1d48aa06295cbc74f975b1 Mon Sep 17 00:00:00 2001 From: Ross Brattain Date: Tue, 29 Nov 2016 13:25:18 -0800 Subject: import new _put_file_shell method from upstream rally upstream openstack rally added new _put_file_* methods we should use these https://github.com/openstack/rally/blob/0.7.0/rally/common/sshutils.py#L270 Updates: imported rally test__put_file_shell unittests quote to prevent word split use -- guard only chmod on cat success Change-Id: I357d1a66b5beddad8042958f4e55d67fc68929f6 Signed-off-by: Ross Brattain --- tests/unit/test_ssh.py | 60 ++++++++++++++++++++++++++++++++++++++++++++++++-- yardstick/ssh.py | 34 ++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_ssh.py b/tests/unit/test_ssh.py index a27052462..083ab039c 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 import mock @@ -162,10 +163,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): @@ -288,6 +289,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 d287b4dac..71dce8102 100644 --- a/yardstick/ssh.py +++ b/yardstick/ssh.py @@ -282,3 +282,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) -- cgit 1.2.3-korg