tests/unit/test_ssh.py

   1 # Copyright 2013: Mirantis Inc.
   2 # All Rights Reserved.
   3 #
   4 #    Licensed under the Apache License, Version 2.0 (the "License"); you may
   5 #    not use this file except in compliance with the License. You may obtain
   6 #    a copy of the License at
   7 #
   8 #         http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 #    Unless required by applicable law or agreed to in writing, software
  11 #    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  12 #    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  13 #    License for the specific language governing permissions and limitations
  14 #    under the License.
  15
  16 # yardstick comment: this file is a modified copy of
  17 # rally/tests/unit/common/test_sshutils.py
  18
  19 from __future__ import absolute_import
  20 import os
  21 import socket
  22 import unittest
  23 from io import StringIO
  24
  25 import mock
  26 from oslo_utils import encodeutils
  27
  28 from yardstick import ssh
  29
  30
  31 class FakeParamikoException(Exception):
  32     pass
  33
  34
  35 class SSHTestCase(unittest.TestCase):
  36     """Test all small SSH methods."""
  37
  38     def setUp(self):
  39         super(SSHTestCase, self).setUp()
  40         self.test_client = ssh.SSH("root", "example.net")
  41
  42     @mock.patch("yardstick.ssh.SSH._get_pkey")
  43     def test_construct(self, mock_ssh__get_pkey):
  44         mock_ssh__get_pkey.return_value = "pkey"
  45         test_ssh = ssh.SSH("root", "example.net", port=33, pkey="key",
  46                            key_filename="kf", password="secret")
  47         mock_ssh__get_pkey.assert_called_once_with("key")
  48         self.assertEqual("root", test_ssh.user)
  49         self.assertEqual("example.net", test_ssh.host)
  50         self.assertEqual(33, test_ssh.port)
  51         self.assertEqual("pkey", test_ssh.pkey)
  52         self.assertEqual("kf", test_ssh.key_filename)
  53         self.assertEqual("secret", test_ssh.password)
  54
  55     def test_construct_default(self):
  56         self.assertEqual("root", self.test_client.user)
  57         self.assertEqual("example.net", self.test_client.host)
  58         self.assertEqual(22, self.test_client.port)
  59         self.assertIsNone(self.test_client.pkey)
  60         self.assertIsNone(self.test_client.key_filename)
  61         self.assertIsNone(self.test_client.password)
  62
  63     @mock.patch("yardstick.ssh.paramiko")
  64     def test__get_pkey_invalid(self, mock_paramiko):
  65         mock_paramiko.SSHException = FakeParamikoException
  66         rsa = mock_paramiko.rsakey.RSAKey
  67         dss = mock_paramiko.dsskey.DSSKey
  68         rsa.from_private_key.side_effect = mock_paramiko.SSHException
  69         dss.from_private_key.side_effect = mock_paramiko.SSHException
  70         self.assertRaises(ssh.SSHError, self.test_client._get_pkey, "key")
  71
  72     @mock.patch("yardstick.ssh.six.moves.StringIO")
  73     @mock.patch("yardstick.ssh.paramiko")
  74     def test__get_pkey_dss(self, mock_paramiko, mock_string_io):
  75         mock_paramiko.SSHException = FakeParamikoException
  76         mock_string_io.return_value = "string_key"
  77         mock_paramiko.dsskey.DSSKey.from_private_key.return_value = "dss_key"
  78         rsa = mock_paramiko.rsakey.RSAKey
  79         rsa.from_private_key.side_effect = mock_paramiko.SSHException
  80         key = self.test_client._get_pkey("key")
  81         dss_calls = mock_paramiko.dsskey.DSSKey.from_private_key.mock_calls
  82         self.assertEqual([mock.call("string_key")], dss_calls)
  83         self.assertEqual(key, "dss_key")
  84         mock_string_io.assert_called_once_with("key")
  85
  86     @mock.patch("yardstick.ssh.six.moves.StringIO")
  87     @mock.patch("yardstick.ssh.paramiko")
  88     def test__get_pkey_rsa(self, mock_paramiko, mock_string_io):
  89         mock_paramiko.SSHException = FakeParamikoException
  90         mock_string_io.return_value = "string_key"
  91         mock_paramiko.rsakey.RSAKey.from_private_key.return_value = "rsa_key"
  92         dss = mock_paramiko.dsskey.DSSKey
  93         dss.from_private_key.side_effect = mock_paramiko.SSHException
  94         key = self.test_client._get_pkey("key")
  95         rsa_calls = mock_paramiko.rsakey.RSAKey.from_private_key.mock_calls
  96         self.assertEqual([mock.call("string_key")], rsa_calls)
  97         self.assertEqual(key, "rsa_key")
  98         mock_string_io.assert_called_once_with("key")
  99
 100     @mock.patch("yardstick.ssh.SSH._get_pkey")
 101     @mock.patch("yardstick.ssh.paramiko")
 102     def test__get_client(self, mock_paramiko, mock_ssh__get_pkey):
 103         mock_ssh__get_pkey.return_value = "key"
 104         fake_client = mock.Mock()
 105         mock_paramiko.SSHClient.return_value = fake_client
 106         mock_paramiko.AutoAddPolicy.return_value = "autoadd"
 107
 108         test_ssh = ssh.SSH("admin", "example.net", pkey="key")
 109         client = test_ssh._get_client()
 110
 111         self.assertEqual(fake_client, client)
 112         client_calls = [
 113             mock.call.set_missing_host_key_policy("autoadd"),
 114             mock.call.connect("example.net", username="admin",
 115                               port=22, pkey="key", key_filename=None,
 116                               password=None,
 117                               allow_agent=False, look_for_keys=False,
 118                               timeout=1),
 119         ]
 120         self.assertEqual(client_calls, client.mock_calls)
 121
 122     def test_close(self):
 123         with mock.patch.object(self.test_client, "_client") as m_client:
 124             self.test_client.close()
 125         m_client.close.assert_called_once_with()
 126         self.assertFalse(self.test_client._client)
 127
 128     @mock.patch("yardstick.ssh.six.moves.StringIO")
 129     def test_execute(self, mock_string_io):
 130         mock_string_io.side_effect = stdio = [mock.Mock(), mock.Mock()]
 131         stdio[0].read.return_value = "stdout fake data"
 132         stdio[1].read.return_value = "stderr fake data"
 133         with mock.patch.object(self.test_client, "run", return_value=0)\
 134                 as mock_run:
 135             status, stdout, stderr = self.test_client.execute(
 136                 "cmd",
 137                 stdin="fake_stdin",
 138                 timeout=43)
 139         mock_run.assert_called_once_with(
 140             "cmd", stdin="fake_stdin", stdout=stdio[0],
 141             stderr=stdio[1], timeout=43, raise_on_error=False)
 142         self.assertEqual(0, status)
 143         self.assertEqual("stdout fake data", stdout)
 144         self.assertEqual("stderr fake data", stderr)
 145
 146     @mock.patch("yardstick.ssh.time")
 147     def test_wait_timeout(self, mock_time):
 148         mock_time.time.side_effect = [1, 50, 150]
 149         self.test_client.execute = mock.Mock(side_effect=[ssh.SSHError,
 150                                                           ssh.SSHError,
 151                                                           0])
 152         self.assertRaises(ssh.SSHTimeout, self.test_client.wait)
 153         self.assertEqual([mock.call("uname")] * 2,
 154                          self.test_client.execute.mock_calls)
 155
 156     @mock.patch("yardstick.ssh.time")
 157     def test_wait(self, mock_time):
 158         mock_time.time.side_effect = [1, 50, 100]
 159         self.test_client.execute = mock.Mock(side_effect=[ssh.SSHError,
 160                                                           ssh.SSHError,
 161                                                           0])
 162         self.test_client.wait()
 163         self.assertEqual([mock.call("uname")] * 3,
 164                          self.test_client.execute.mock_calls)
 165
 166     @mock.patch("yardstick.ssh.paramiko")
 167     def test_send_command(self, mock_paramiko):
 168         paramiko_sshclient = self.test_client._get_client()
 169         with mock.patch.object(paramiko_sshclient, "exec_command") \
 170                 as mock_paramiko_exec_command:
 171             self.test_client.send_command('cmd')
 172         mock_paramiko_exec_command.assert_called_once_with('cmd',
 173                                                            get_pty=True)
 174
 175
 176 class SSHRunTestCase(unittest.TestCase):
 177     """Test SSH.run method in different aspects.
 178
 179     Also tested method "execute".
 180     """
 181
 182     def setUp(self):
 183         super(SSHRunTestCase, self).setUp()
 184
 185         self.fake_client = mock.Mock()
 186         self.fake_session = mock.Mock()
 187         self.fake_transport = mock.Mock()
 188
 189         self.fake_transport.open_session.return_value = self.fake_session
 190         self.fake_client.get_transport.return_value = self.fake_transport
 191
 192         self.fake_session.recv_ready.return_value = False
 193         self.fake_session.recv_stderr_ready.return_value = False
 194         self.fake_session.send_ready.return_value = False
 195         self.fake_session.exit_status_ready.return_value = True
 196         self.fake_session.recv_exit_status.return_value = 0
 197
 198         self.test_client = ssh.SSH("admin", "example.net")
 199         self.test_client._get_client = mock.Mock(return_value=self.fake_client)
 200
 201     @mock.patch("yardstick.ssh.select")
 202     def test_execute(self, mock_select):
 203         mock_select.select.return_value = ([], [], [])
 204         self.fake_session.recv_ready.side_effect = [1, 0, 0]
 205         self.fake_session.recv_stderr_ready.side_effect = [1, 0]
 206         self.fake_session.recv.return_value = "ok"
 207         self.fake_session.recv_stderr.return_value = "error"
 208         self.fake_session.exit_status_ready.return_value = 1
 209         self.fake_session.recv_exit_status.return_value = 127
 210         self.assertEqual((127, "ok", "error"), self.test_client.execute("cmd"))
 211         self.fake_session.exec_command.assert_called_once_with("cmd")
 212
 213     @mock.patch("yardstick.ssh.select")
 214     def test_execute_args(self, mock_select):
 215         mock_select.select.return_value = ([], [], [])
 216         self.fake_session.recv_ready.side_effect = [1, 0, 0]
 217         self.fake_session.recv_stderr_ready.side_effect = [1, 0]
 218         self.fake_session.recv.return_value = "ok"
 219         self.fake_session.recv_stderr.return_value = "error"
 220         self.fake_session.exit_status_ready.return_value = 1
 221         self.fake_session.recv_exit_status.return_value = 127
 222
 223         result = self.test_client.execute("cmd arg1 'arg2 with space'")
 224         self.assertEqual((127, "ok", "error"), result)
 225         self.fake_session.exec_command.assert_called_once_with(
 226             "cmd arg1 'arg2 with space'")
 227
 228     @mock.patch("yardstick.ssh.select")
 229     def test_run(self, mock_select):
 230         mock_select.select.return_value = ([], [], [])
 231         self.assertEqual(0, self.test_client.run("cmd"))
 232
 233     @mock.patch("yardstick.ssh.select")
 234     def test_run_nonzero_status(self, mock_select):
 235         mock_select.select.return_value = ([], [], [])
 236         self.fake_session.recv_exit_status.return_value = 1
 237         self.assertRaises(ssh.SSHError, self.test_client.run, "cmd")
 238         self.assertEqual(1, self.test_client.run("cmd", raise_on_error=False))
 239
 240     @mock.patch("yardstick.ssh.select")
 241     def test_run_stdout(self, mock_select):
 242         mock_select.select.return_value = ([], [], [])
 243         self.fake_session.recv_ready.side_effect = [True, True, False]
 244         self.fake_session.recv.side_effect = ["ok1", "ok2"]
 245         stdout = mock.Mock()
 246         self.test_client.run("cmd", stdout=stdout)
 247         self.assertEqual([mock.call("ok1"), mock.call("ok2")],
 248                          stdout.write.mock_calls)
 249
 250     @mock.patch("yardstick.ssh.select")
 251     def test_run_stderr(self, mock_select):
 252         mock_select.select.return_value = ([], [], [])
 253         self.fake_session.recv_stderr_ready.side_effect = [True, False]
 254         self.fake_session.recv_stderr.return_value = "error"
 255         stderr = mock.Mock()
 256         self.test_client.run("cmd", stderr=stderr)
 257         stderr.write.assert_called_once_with("error")
 258
 259     @mock.patch("yardstick.ssh.select")
 260     def test_run_stdin(self, mock_select):
 261         """Test run method with stdin.
 262
 263         Third send call was called with "e2" because only 3 bytes was sent
 264         by second call. So remainig 2 bytes of "line2" was sent by third call.
 265         """
 266         mock_select.select.return_value = ([], [], [])
 267         self.fake_session.exit_status_ready.side_effect = [0, 0, 0, True]
 268         self.fake_session.send_ready.return_value = True
 269         self.fake_session.send.side_effect = [5, 3, 2]
 270         fake_stdin = mock.Mock()
 271         fake_stdin.read.side_effect = ["line1", "line2", ""]
 272         fake_stdin.closed = False
 273
 274         def close():
 275             fake_stdin.closed = True
 276         fake_stdin.close = mock.Mock(side_effect=close)
 277         self.test_client.run("cmd", stdin=fake_stdin)
 278         call = mock.call
 279         send_calls = [call(encodeutils.safe_encode("line1", "utf-8")),
 280                       call(encodeutils.safe_encode("line2", "utf-8")),
 281                       call(encodeutils.safe_encode("e2", "utf-8"))]
 282         self.assertEqual(send_calls, self.fake_session.send.mock_calls)
 283
 284     @mock.patch("yardstick.ssh.select")
 285     def test_run_stdin_keep_open(self, mock_select):
 286         """Test run method with stdin.
 287
 288         Third send call was called with "e2" because only 3 bytes was sent
 289         by second call. So remainig 2 bytes of "line2" was sent by third call.
 290         """
 291         mock_select.select.return_value = ([], [], [])
 292         self.fake_session.exit_status_ready.side_effect = [0, 0, 0, True]
 293         self.fake_session.send_ready.return_value = True
 294         self.fake_session.send.side_effect = len
 295         fake_stdin = StringIO(u"line1\nline2\n")
 296         self.test_client.run("cmd", stdin=fake_stdin, keep_stdin_open=True)
 297         call = mock.call
 298         send_calls = [call(encodeutils.safe_encode("line1\nline2\n", "utf-8"))]
 299         self.assertEqual(send_calls, self.fake_session.send.mock_calls)
 300
 301     @mock.patch("yardstick.ssh.select")
 302     def test_run_select_error(self, mock_select):
 303         self.fake_session.exit_status_ready.return_value = False
 304         mock_select.select.return_value = ([], [], [True])
 305         self.assertRaises(ssh.SSHError, self.test_client.run, "cmd")
 306
 307     @mock.patch("yardstick.ssh.time")
 308     @mock.patch("yardstick.ssh.select")
 309     def test_run_timemout(self, mock_select, mock_time):
 310         mock_time.time.side_effect = [1, 3700]
 311         mock_select.select.return_value = ([], [], [])
 312         self.fake_session.exit_status_ready.return_value = False
 313         self.assertRaises(ssh.SSHTimeout, self.test_client.run, "cmd")
 314
 315     @mock.patch("yardstick.ssh.open", create=True)
 316     def test__put_file_shell(self, mock_open):
 317         with mock.patch.object(self.test_client, "run") as run_mock:
 318             self.test_client._put_file_shell("localfile", "remotefile", 0o42)
 319             run_mock.assert_called_once_with(
 320                 'cat > "remotefile"&& chmod -- 042 "remotefile"',
 321                 stdin=mock_open.return_value.__enter__.return_value)
 322
 323     @mock.patch("yardstick.ssh.open", create=True)
 324     def test__put_file_shell_space(self, mock_open):
 325         with mock.patch.object(self.test_client, "run") as run_mock:
 326             self.test_client._put_file_shell("localfile",
 327                                              "filename with space", 0o42)
 328             run_mock.assert_called_once_with(
 329                 'cat > "filename with space"&& chmod -- 042 "filename with '
 330                 'space"',
 331                 stdin=mock_open.return_value.__enter__.return_value)
 332
 333     @mock.patch("yardstick.ssh.open", create=True)
 334     def test__put_file_shell_tilde(self, mock_open):
 335         with mock.patch.object(self.test_client, "run") as run_mock:
 336             self.test_client._put_file_shell("localfile", "~/remotefile", 0o42)
 337             run_mock.assert_called_once_with(
 338                 'cat > ~/"remotefile"&& chmod -- 042 ~/"remotefile"',
 339                 stdin=mock_open.return_value.__enter__.return_value)
 340
 341     @mock.patch("yardstick.ssh.open", create=True)
 342     def test__put_file_shell_tilde_spaces(self, mock_open):
 343         with mock.patch.object(self.test_client, "run") as run_mock:
 344             self.test_client._put_file_shell("localfile", "~/file with space",
 345                                              0o42)
 346             run_mock.assert_called_once_with(
 347                 'cat > ~/"file with space"&& chmod -- 042 ~/"file with space"',
 348                 stdin=mock_open.return_value.__enter__.return_value)
 349
 350     @mock.patch("yardstick.ssh.os.stat")
 351     def test__put_file_sftp(self, mock_stat):
 352         sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
 353         sftp.__enter__.return_value = sftp
 354
 355         mock_stat.return_value = os.stat_result([0o753] + [0] * 9)
 356
 357         self.test_client._put_file_sftp("localfile", "remotefile")
 358
 359         sftp.put.assert_called_once_with("localfile", "remotefile")
 360         mock_stat.assert_called_once_with("localfile")
 361         sftp.chmod.assert_called_once_with("remotefile", 0o753)
 362         sftp.__exit__.assert_called_once_with(None, None, None)
 363
 364     def test__put_file_sftp_mode(self):
 365         sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
 366         sftp.__enter__.return_value = sftp
 367
 368         self.test_client._put_file_sftp("localfile", "remotefile", mode=0o753)
 369
 370         sftp.put.assert_called_once_with("localfile", "remotefile")
 371         sftp.chmod.assert_called_once_with("remotefile", 0o753)
 372         sftp.__exit__.assert_called_once_with(None, None, None)
 373
 374     def test_put_file_SSHException(self):
 375         exc = ssh.paramiko.SSHException
 376         self.test_client._put_file_sftp = mock.Mock(side_effect=exc())
 377         self.test_client._put_file_shell = mock.Mock()
 378
 379         self.test_client.put_file("foo", "bar", 42)
 380         self.test_client._put_file_sftp.assert_called_once_with("foo", "bar",
 381                                                                 mode=42)
 382         self.test_client._put_file_shell.assert_called_once_with("foo", "bar",
 383                                                                  mode=42)
 384
 385     def test_put_file_socket_error(self):
 386         exc = socket.error
 387         self.test_client._put_file_sftp = mock.Mock(side_effect=exc())
 388         self.test_client._put_file_shell = mock.Mock()
 389
 390         self.test_client.put_file("foo", "bar", 42)
 391         self.test_client._put_file_sftp.assert_called_once_with("foo", "bar",
 392                                                                 mode=42)
 393         self.test_client._put_file_shell.assert_called_once_with("foo", "bar",
 394                                                                  mode=42)
 395
 396
 397 def main():
 398     unittest.main()
 399
 400
 401 if __name__ == '__main__':
 402     main()