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