yardstick/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 import os
  17 import socket
  18 import unittest
  19 from io import StringIO
  20 from itertools import count
  21
  22 import mock
  23 from oslo_utils import encodeutils
  24
  25 from yardstick.common import exceptions
  26 from yardstick import ssh
  27 from yardstick.ssh import SSH
  28 from yardstick.ssh import AutoConnectSSH
  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     @mock.patch("yardstick.ssh.SSH._get_pkey")
  56     def test_ssh_from_node(self, mock_ssh__get_pkey):
  57         mock_ssh__get_pkey.return_value = "pkey"
  58         node = {
  59             "user": "root", "ip": "example.net", "ssh_port": 33,
  60             "key_filename": "kf", "password": "secret"
  61         }
  62         test_ssh = ssh.SSH.from_node(node)
  63         self.assertEqual("root", test_ssh.user)
  64         self.assertEqual("example.net", test_ssh.host)
  65         self.assertEqual(33, test_ssh.port)
  66         self.assertEqual("kf", test_ssh.key_filename)
  67         self.assertEqual("secret", test_ssh.password)
  68
  69     @mock.patch("yardstick.ssh.SSH._get_pkey")
  70     def test_ssh_from_node_password_default(self, mock_ssh__get_pkey):
  71         mock_ssh__get_pkey.return_value = "pkey"
  72         node = {
  73             "user": "root", "ip": "example.net", "ssh_port": 33,
  74             "key_filename": "kf"
  75         }
  76         test_ssh = ssh.SSH.from_node(node)
  77         self.assertEqual("root", test_ssh.user)
  78         self.assertEqual("example.net", test_ssh.host)
  79         self.assertEqual(33, test_ssh.port)
  80         self.assertEqual("kf", test_ssh.key_filename)
  81         self.assertIsNone(test_ssh.password)
  82
  83     @mock.patch("yardstick.ssh.SSH._get_pkey")
  84     def test_ssh_from_node_ssh_port_default(self, mock_ssh__get_pkey):
  85         mock_ssh__get_pkey.return_value = "pkey"
  86         node = {
  87             "user": "root", "ip": "example.net",
  88             "key_filename": "kf", "password": "secret"
  89         }
  90         test_ssh = ssh.SSH.from_node(node)
  91         self.assertEqual("root", test_ssh.user)
  92         self.assertEqual("example.net", test_ssh.host)
  93         self.assertEqual(ssh.SSH.SSH_PORT, test_ssh.port)
  94         self.assertEqual("kf", test_ssh.key_filename)
  95         self.assertEqual("secret", test_ssh.password)
  96
  97     @mock.patch("yardstick.ssh.SSH._get_pkey")
  98     def test_ssh_from_node_key_filename_default(self, mock_ssh__get_pkey):
  99         mock_ssh__get_pkey.return_value = "pkey"
 100         node = {
 101             "user": "root", "ip": "example.net", "ssh_port": 33,
 102             "password": "secret"
 103         }
 104         test_ssh = ssh.SSH.from_node(node)
 105         self.assertEqual("root", test_ssh.user)
 106         self.assertEqual("example.net", test_ssh.host)
 107         self.assertEqual(33, test_ssh.port)
 108         self.assertIsNone(test_ssh.key_filename)
 109         self.assertEqual("secret", test_ssh.password)
 110
 111     def test_construct_default(self):
 112         self.assertEqual("root", self.test_client.user)
 113         self.assertEqual("example.net", self.test_client.host)
 114         self.assertEqual(22, self.test_client.port)
 115         self.assertIsNone(self.test_client.pkey)
 116         self.assertIsNone(self.test_client.key_filename)
 117         self.assertIsNone(self.test_client.password)
 118
 119     @mock.patch("yardstick.ssh.paramiko")
 120     def test__get_pkey_invalid(self, mock_paramiko):
 121         mock_paramiko.SSHException = FakeParamikoException
 122         rsa = mock_paramiko.rsakey.RSAKey
 123         dss = mock_paramiko.dsskey.DSSKey
 124         rsa.from_private_key.side_effect = mock_paramiko.SSHException
 125         dss.from_private_key.side_effect = mock_paramiko.SSHException
 126         self.assertRaises(exceptions.SSHError, self.test_client._get_pkey, "key")
 127
 128     @mock.patch("yardstick.ssh.six.moves.StringIO")
 129     @mock.patch("yardstick.ssh.paramiko")
 130     def test__get_pkey_dss(self, mock_paramiko, mock_string_io):
 131         mock_paramiko.SSHException = FakeParamikoException
 132         mock_string_io.return_value = "string_key"
 133         mock_paramiko.dsskey.DSSKey.from_private_key.return_value = "dss_key"
 134         rsa = mock_paramiko.rsakey.RSAKey
 135         rsa.from_private_key.side_effect = mock_paramiko.SSHException
 136         key = self.test_client._get_pkey("key")
 137         dss_calls = mock_paramiko.dsskey.DSSKey.from_private_key.mock_calls
 138         self.assertEqual([mock.call("string_key")], dss_calls)
 139         self.assertEqual(key, "dss_key")
 140         mock_string_io.assert_called_once_with("key")
 141
 142     @mock.patch("yardstick.ssh.six.moves.StringIO")
 143     @mock.patch("yardstick.ssh.paramiko")
 144     def test__get_pkey_rsa(self, mock_paramiko, mock_string_io):
 145         mock_paramiko.SSHException = FakeParamikoException
 146         mock_string_io.return_value = "string_key"
 147         mock_paramiko.rsakey.RSAKey.from_private_key.return_value = "rsa_key"
 148         dss = mock_paramiko.dsskey.DSSKey
 149         dss.from_private_key.side_effect = mock_paramiko.SSHException
 150         key = self.test_client._get_pkey("key")
 151         rsa_calls = mock_paramiko.rsakey.RSAKey.from_private_key.mock_calls
 152         self.assertEqual([mock.call("string_key")], rsa_calls)
 153         self.assertEqual(key, "rsa_key")
 154         mock_string_io.assert_called_once_with("key")
 155
 156     @mock.patch("yardstick.ssh.SSH._get_pkey")
 157     @mock.patch("yardstick.ssh.paramiko")
 158     def test__get_client(self, mock_paramiko, mock_ssh__get_pkey):
 159         mock_ssh__get_pkey.return_value = "key"
 160         fake_client = mock.Mock()
 161         mock_paramiko.SSHClient.return_value = fake_client
 162         mock_paramiko.AutoAddPolicy.return_value = "autoadd"
 163
 164         test_ssh = ssh.SSH("admin", "example.net", pkey="key")
 165         client = test_ssh._get_client()
 166
 167         self.assertEqual(fake_client, client)
 168         client_calls = [
 169             mock.call.set_missing_host_key_policy("autoadd"),
 170             mock.call.connect("example.net", username="admin",
 171                               port=22, pkey="key", key_filename=None,
 172                               password=None,
 173                               allow_agent=False, look_for_keys=False,
 174                               timeout=1),
 175         ]
 176         self.assertEqual(client_calls, client.mock_calls)
 177
 178     @mock.patch("yardstick.ssh.SSH._get_pkey")
 179     @mock.patch("yardstick.ssh.paramiko")
 180     def test__get_client_with_exception(self, mock_paramiko, mock_ssh__get_pkey):
 181         class MyError(Exception):
 182             pass
 183
 184         mock_ssh__get_pkey.return_value = "pkey"
 185         fake_client = mock.Mock()
 186         fake_client.connect.side_effect = MyError
 187         fake_client.set_missing_host_key_policy.return_value = None
 188         mock_paramiko.SSHClient.return_value = fake_client
 189         mock_paramiko.AutoAddPolicy.return_value = "autoadd"
 190
 191         test_ssh = ssh.SSH("admin", "example.net", pkey="key")
 192
 193         with self.assertRaises(exceptions.SSHError) as raised:
 194             test_ssh._get_client()
 195
 196         mock_paramiko.SSHClient.assert_called_once()
 197         mock_paramiko.AutoAddPolicy.assert_called_once()
 198         fake_client.set_missing_host_key_policy.assert_called_once()
 199         fake_client.connect.assert_called_once()
 200         exc_str = str(raised.exception)
 201         self.assertIn('raised during connect', exc_str)
 202         self.assertIn('MyError', exc_str)
 203
 204     @mock.patch("yardstick.ssh.SSH._get_pkey")
 205     @mock.patch("yardstick.ssh.paramiko")
 206     def test_copy(self, mock_paramiko, mock_ssh__get_pkey):
 207         mock_ssh__get_pkey.return_value = "pkey"
 208         fake_client = mock.Mock()
 209         fake_client.connect.side_effect = IOError
 210         mock_paramiko.SSHClient.return_value = fake_client
 211         mock_paramiko.AutoAddPolicy.return_value = "autoadd"
 212
 213         test_ssh = ssh.SSH("admin", "example.net", pkey="key")
 214         result = test_ssh.copy()
 215         self.assertIsNot(test_ssh, result)
 216
 217     def test_close(self):
 218         with mock.patch.object(self.test_client, "_client") as m_client:
 219             self.test_client.close()
 220         m_client.close.assert_called_once_with()
 221         self.assertFalse(self.test_client._client)
 222
 223     @mock.patch("yardstick.ssh.six.moves.StringIO")
 224     def test_execute(self, mock_string_io):
 225         mock_string_io.side_effect = stdio = [mock.Mock(), mock.Mock()]
 226         stdio[0].read.return_value = "stdout fake data"
 227         stdio[1].read.return_value = "stderr fake data"
 228         with mock.patch.object(self.test_client, "run", return_value=0)\
 229                 as mock_run:
 230             status, stdout, stderr = self.test_client.execute(
 231                 "cmd",
 232                 stdin="fake_stdin",
 233                 timeout=43)
 234         mock_run.assert_called_once_with(
 235             "cmd", stdin="fake_stdin", stdout=stdio[0],
 236             stderr=stdio[1], timeout=43, raise_on_error=False)
 237         self.assertEqual(0, status)
 238         self.assertEqual("stdout fake data", stdout)
 239         self.assertEqual("stderr fake data", stderr)
 240
 241     @mock.patch("yardstick.ssh.six.moves.StringIO")
 242     def test_execute_raise_on_error_passed(self, mock_string_io):
 243         mock_string_io.side_effect = stdio = [mock.Mock(), mock.Mock()]
 244         stdio[0].read.return_value = "stdout fake data"
 245         stdio[1].read.return_value = "stderr fake data"
 246         with mock.patch.object(self.test_client, "run", return_value=0) \
 247                 as mock_run:
 248             status, stdout, stderr = self.test_client.execute(
 249                 "cmd",
 250                 stdin="fake_stdin",
 251                 timeout=43,
 252                 raise_on_error=True)
 253         mock_run.assert_called_once_with(
 254             "cmd", stdin="fake_stdin", stdout=stdio[0],
 255             stderr=stdio[1], timeout=43, raise_on_error=True)
 256         self.assertEqual(0, status)
 257         self.assertEqual("stdout fake data", stdout)
 258         self.assertEqual("stderr fake data", stderr)
 259
 260     @mock.patch("yardstick.ssh.time")
 261     def test_wait_timeout(self, mock_time):
 262         mock_time.time.side_effect = [1, 50, 150]
 263         self.test_client.execute = mock.Mock(side_effect=[exceptions.SSHError,
 264                                                           exceptions.SSHError,
 265                                                           0])
 266         self.assertRaises(exceptions.SSHTimeout, self.test_client.wait)
 267         self.assertEqual([mock.call("uname")] * 2,
 268                          self.test_client.execute.mock_calls)
 269
 270     @mock.patch("yardstick.ssh.time")
 271     def test_wait(self, mock_time):
 272         mock_time.time.side_effect = [1, 50, 100]
 273         self.test_client.execute = mock.Mock(side_effect=[exceptions.SSHError,
 274                                                           exceptions.SSHError,
 275                                                           0])
 276         self.test_client.wait()
 277         self.assertEqual([mock.call("uname")] * 3,
 278                          self.test_client.execute.mock_calls)
 279
 280     @mock.patch("yardstick.ssh.paramiko")
 281     def test_send_command(self, _):
 282         paramiko_sshclient = self.test_client._get_client()
 283         with mock.patch.object(paramiko_sshclient, "exec_command") \
 284                 as mock_paramiko_exec_command:
 285             self.test_client.send_command('cmd')
 286         mock_paramiko_exec_command.assert_called_once_with('cmd',
 287                                                            get_pty=True)
 288
 289     @mock.patch("yardstick.ssh.paramiko")
 290     def test_interactive_terminal_open(self, mock_paramiko):
 291         fake_client = mock.Mock()
 292         fake_session = mock.Mock()
 293         fake_session.recv.return_value = ":~# "
 294         fake_transport = mock.Mock()
 295         fake_transport.open_session.return_value = fake_session
 296         fake_client.get_transport.return_value = fake_transport
 297         mock_paramiko.SSHClient.return_value = fake_client
 298
 299         test_ssh = ssh.SSH("admin", "example.net", pkey="key")
 300         result = test_ssh.interactive_terminal_open()
 301         self.assertEqual(fake_session, result)
 302
 303     @mock.patch("yardstick.ssh.paramiko")
 304     def test_interactive_terminal_exec_command(self, mock_paramiko):
 305         fake_client = mock.Mock()
 306         fake_session = mock.Mock()
 307         fake_session.recv.return_value = "stdout fake data"
 308         fake_transport = mock.Mock()
 309         fake_transport.open_session.return_value = fake_session
 310         fake_client.get_transport.return_value = fake_transport
 311         mock_paramiko.SSHClient.return_value = fake_client
 312
 313         test_ssh = ssh.SSH("admin", "example.net", pkey="key")
 314         with mock.patch.object(fake_session, "sendall") \
 315                 as mock_paramiko_send_command:
 316             result = test_ssh.interactive_terminal_exec_command(fake_session,
 317                                                                 'cmd', "vat# ")
 318         self.assertEqual("stdout fake data", result)
 319         mock_paramiko_send_command.assert_called_once_with('cmd\n')
 320
 321     @mock.patch("yardstick.ssh.paramiko")
 322     def test_interactive_terminal_close(self, _):
 323         fake_session = mock.Mock()
 324         paramiko_sshclient = self.test_client._get_client()
 325         paramiko_sshclient.get_transport.open_session.return_value = fake_session
 326         with mock.patch.object(fake_session, "close") \
 327                 as mock_paramiko_terminal_close:
 328             self.test_client.interactive_terminal_close(fake_session)
 329         mock_paramiko_terminal_close.assert_called_once_with()
 330
 331
 332 class SSHRunTestCase(unittest.TestCase):
 333     """Test SSH.run method in different aspects.
 334
 335     Also tested method "execute".
 336     """
 337
 338     def setUp(self):
 339         super(SSHRunTestCase, self).setUp()
 340
 341         self.fake_client = mock.Mock()
 342         self.fake_session = mock.Mock()
 343         self.fake_transport = mock.Mock()
 344
 345         self.fake_transport.open_session.return_value = self.fake_session
 346         self.fake_client.get_transport.return_value = self.fake_transport
 347
 348         self.fake_session.recv_ready.return_value = False
 349         self.fake_session.recv_stderr_ready.return_value = False
 350         self.fake_session.send_ready.return_value = False
 351         self.fake_session.exit_status_ready.return_value = True
 352         self.fake_session.recv_exit_status.return_value = 0
 353
 354         self.test_client = ssh.SSH("admin", "example.net")
 355         self.test_client._get_client = mock.Mock(return_value=self.fake_client)
 356
 357     @mock.patch("yardstick.ssh.select")
 358     def test_execute(self, mock_select):
 359         mock_select.select.return_value = ([], [], [])
 360         self.fake_session.recv_ready.side_effect = [1, 0, 0]
 361         self.fake_session.recv_stderr_ready.side_effect = [1, 0]
 362         self.fake_session.recv.return_value = "ok"
 363         self.fake_session.recv_stderr.return_value = "error"
 364         self.fake_session.exit_status_ready.return_value = 1
 365         self.fake_session.recv_exit_status.return_value = 127
 366         self.assertEqual((127, "ok", "error"), self.test_client.execute("cmd"))
 367         self.fake_session.exec_command.assert_called_once_with("cmd")
 368
 369     @mock.patch("yardstick.ssh.select")
 370     def test_execute_args(self, mock_select):
 371         mock_select.select.return_value = ([], [], [])
 372         self.fake_session.recv_ready.side_effect = [1, 0, 0]
 373         self.fake_session.recv_stderr_ready.side_effect = [1, 0]
 374         self.fake_session.recv.return_value = "ok"
 375         self.fake_session.recv_stderr.return_value = "error"
 376         self.fake_session.exit_status_ready.return_value = 1
 377         self.fake_session.recv_exit_status.return_value = 127
 378
 379         result = self.test_client.execute("cmd arg1 'arg2 with space'")
 380         self.assertEqual((127, "ok", "error"), result)
 381         self.fake_session.exec_command.assert_called_once_with(
 382             "cmd arg1 'arg2 with space'")
 383
 384     @mock.patch("yardstick.ssh.select")
 385     def test_run(self, mock_select):
 386         mock_select.select.return_value = ([], [], [])
 387         self.assertEqual(0, self.test_client.run("cmd"))
 388
 389     @mock.patch("yardstick.ssh.select")
 390     def test_run_nonzero_status(self, mock_select):
 391         mock_select.select.return_value = ([], [], [])
 392         self.fake_session.recv_exit_status.return_value = 1
 393         self.assertRaises(exceptions.SSHError, self.test_client.run, "cmd")
 394         self.assertEqual(1, self.test_client.run("cmd", raise_on_error=False))
 395
 396     @mock.patch("yardstick.ssh.select")
 397     def test_run_stdout(self, mock_select):
 398         mock_select.select.return_value = ([], [], [])
 399         self.fake_session.recv_ready.side_effect = [True, True, False]
 400         self.fake_session.recv.side_effect = ["ok1", "ok2"]
 401         stdout = mock.Mock()
 402         self.test_client.run("cmd", stdout=stdout)
 403         self.assertEqual([mock.call("ok1"), mock.call("ok2")],
 404                          stdout.write.mock_calls)
 405
 406     @mock.patch("yardstick.ssh.select")
 407     def test_run_stderr(self, mock_select):
 408         mock_select.select.return_value = ([], [], [])
 409         self.fake_session.recv_stderr_ready.side_effect = [True, False]
 410         self.fake_session.recv_stderr.return_value = "error"
 411         stderr = mock.Mock()
 412         self.test_client.run("cmd", stderr=stderr)
 413         stderr.write.assert_called_once_with("error")
 414
 415     @mock.patch("yardstick.ssh.select")
 416     def test_run_stdin(self, mock_select):
 417         """Test run method with stdin.
 418
 419         Third send call was called with "e2" because only 3 bytes was sent
 420         by second call. So remainig 2 bytes of "line2" was sent by third call.
 421         """
 422         mock_select.select.return_value = ([], [], [])
 423         self.fake_session.exit_status_ready.side_effect = [0, 0, 0, True]
 424         self.fake_session.send_ready.return_value = True
 425         self.fake_session.send.side_effect = [5, 3, 2]
 426         fake_stdin = mock.Mock()
 427         fake_stdin.read.side_effect = ["line1", "line2", ""]
 428         fake_stdin.closed = False
 429
 430         def close():
 431             fake_stdin.closed = True
 432         fake_stdin.close = mock.Mock(side_effect=close)
 433         self.test_client.run("cmd", stdin=fake_stdin)
 434         call = mock.call
 435         send_calls = [call(encodeutils.safe_encode("line1", "utf-8")),
 436                       call(encodeutils.safe_encode("line2", "utf-8")),
 437                       call(encodeutils.safe_encode("e2", "utf-8"))]
 438         self.assertEqual(send_calls, self.fake_session.send.mock_calls)
 439
 440     @mock.patch("yardstick.ssh.select")
 441     def test_run_stdin_keep_open(self, mock_select):
 442         """Test run method with stdin.
 443
 444         Third send call was called with "e2" because only 3 bytes was sent
 445         by second call. So remainig 2 bytes of "line2" was sent by third call.
 446         """
 447         mock_select.select.return_value = ([], [], [])
 448         self.fake_session.exit_status_ready.side_effect = [0, 0, 0, True]
 449         self.fake_session.send_ready.return_value = True
 450         self.fake_session.send.side_effect = len
 451         fake_stdin = StringIO(u"line1\nline2\n")
 452         self.test_client.run("cmd", stdin=fake_stdin, keep_stdin_open=True)
 453         call = mock.call
 454         send_calls = [call(encodeutils.safe_encode("line1\nline2\n", "utf-8"))]
 455         self.assertEqual(send_calls, self.fake_session.send.mock_calls)
 456
 457     @mock.patch("yardstick.ssh.select")
 458     def test_run_select_error(self, mock_select):
 459         self.fake_session.exit_status_ready.return_value = False
 460         mock_select.select.return_value = ([], [], [True])
 461         self.assertRaises(exceptions.SSHError, self.test_client.run, "cmd")
 462
 463     @mock.patch("yardstick.ssh.time")
 464     @mock.patch("yardstick.ssh.select")
 465     def test_run_timemout(self, mock_select, mock_time):
 466         mock_time.time.side_effect = [1, 3700]
 467         mock_select.select.return_value = ([], [], [])
 468         self.fake_session.exit_status_ready.return_value = False
 469         self.assertRaises(exceptions.SSHTimeout, self.test_client.run, "cmd")
 470
 471     @mock.patch("yardstick.ssh.open", create=True)
 472     def test__put_file_shell(self, mock_open):
 473         with mock.patch.object(self.test_client, "run") as run_mock:
 474             self.test_client._put_file_shell("localfile", "remotefile", 0o42)
 475             run_mock.assert_called_once_with(
 476                 'cat > "remotefile"&& chmod -- 042 "remotefile"',
 477                 stdin=mock_open.return_value.__enter__.return_value)
 478
 479     @mock.patch("yardstick.ssh.open", create=True)
 480     def test__put_file_shell_space(self, mock_open):
 481         with mock.patch.object(self.test_client, "run") as run_mock:
 482             self.test_client._put_file_shell("localfile",
 483                                              "filename with space", 0o42)
 484             run_mock.assert_called_once_with(
 485                 'cat > "filename with space"&& chmod -- 042 "filename with '
 486                 'space"',
 487                 stdin=mock_open.return_value.__enter__.return_value)
 488
 489     @mock.patch("yardstick.ssh.open", create=True)
 490     def test__put_file_shell_tilde(self, mock_open):
 491         with mock.patch.object(self.test_client, "run") as run_mock:
 492             self.test_client._put_file_shell("localfile", "~/remotefile", 0o42)
 493             run_mock.assert_called_once_with(
 494                 'cat > ~/"remotefile"&& chmod -- 042 ~/"remotefile"',
 495                 stdin=mock_open.return_value.__enter__.return_value)
 496
 497     @mock.patch("yardstick.ssh.open", create=True)
 498     def test__put_file_shell_tilde_spaces(self, mock_open):
 499         with mock.patch.object(self.test_client, "run") as run_mock:
 500             self.test_client._put_file_shell("localfile", "~/file with space",
 501                                              0o42)
 502             run_mock.assert_called_once_with(
 503                 'cat > ~/"file with space"&& chmod -- 042 ~/"file with space"',
 504                 stdin=mock_open.return_value.__enter__.return_value)
 505
 506     @mock.patch("yardstick.ssh.os.stat")
 507     def test__put_file_sftp(self, mock_stat):
 508         sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
 509         sftp.__enter__.return_value = sftp
 510
 511         mock_stat.return_value = os.stat_result([0o753] + [0] * 9)
 512
 513         self.test_client._put_file_sftp("localfile", "remotefile")
 514
 515         sftp.put.assert_called_once_with("localfile", "remotefile")
 516         mock_stat.assert_any_call("localfile")
 517         sftp.chmod.assert_any_call("remotefile", 0o753)
 518         sftp.__exit__.assert_called_once_with(None, None, None)
 519
 520     def test__put_file_sftp_mode(self):
 521         sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
 522         sftp.__enter__.return_value = sftp
 523
 524         self.test_client._put_file_sftp("localfile", "remotefile", mode=0o753)
 525
 526         sftp.put.assert_called_once_with("localfile", "remotefile")
 527         sftp.chmod.assert_called_once_with("remotefile", 0o753)
 528         sftp.__exit__.assert_called_once_with(None, None, None)
 529
 530     def test_put_file_SSHException(self):
 531         exc = ssh.paramiko.SSHException
 532         self.test_client._put_file_sftp = mock.Mock(side_effect=exc())
 533         self.test_client._put_file_shell = mock.Mock()
 534
 535         self.test_client.put_file("foo", "bar", 42)
 536         self.test_client._put_file_sftp.assert_called_once_with("foo", "bar",
 537                                                                 mode=42)
 538         self.test_client._put_file_shell.assert_called_once_with("foo", "bar",
 539                                                                  mode=42)
 540
 541     def test_put_file_socket_error(self):
 542         exc = socket.error
 543         self.test_client._put_file_sftp = mock.Mock(side_effect=exc())
 544         self.test_client._put_file_shell = mock.Mock()
 545
 546         self.test_client.put_file("foo", "bar", 42)
 547         self.test_client._put_file_sftp.assert_called_once_with("foo", "bar",
 548                                                                 mode=42)
 549         self.test_client._put_file_shell.assert_called_once_with("foo", "bar",
 550                                                                  mode=42)
 551
 552     @mock.patch("yardstick.ssh.os.stat")
 553     def test_put_file_obj_with_mode(self, mock_stat):
 554         sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
 555         sftp.__enter__.return_value = sftp
 556
 557         mock_stat.return_value = os.stat_result([0o753] + [0] * 9)
 558
 559         self.test_client.put_file_obj("localfile", "remotefile", 'my_mode')
 560
 561         sftp.__enter__.assert_called_once()
 562         sftp.putfo.assert_called_once_with("localfile", "remotefile")
 563         sftp.chmod.assert_called_once_with("remotefile", 'my_mode')
 564         sftp.__exit__.assert_called_once_with(None, None, None)
 565
 566
 567 class TestAutoConnectSSH(unittest.TestCase):
 568
 569     def test__connect_loop(self):
 570         auto_connect_ssh = AutoConnectSSH('user1', 'host1', wait=0)
 571         auto_connect_ssh._get_client = mock__get_client = mock.Mock()
 572
 573         auto_connect_ssh._connect()
 574         mock__get_client.assert_called_once()
 575
 576     def test___init___negative(self):
 577         with self.assertRaises(TypeError):
 578             AutoConnectSSH('user1', 'host1', wait=['wait'])
 579
 580         with self.assertRaises(ValueError):
 581             AutoConnectSSH('user1', 'host1', wait='wait')
 582
 583     @mock.patch('yardstick.ssh.time')
 584     def test__connect_loop_ssh_error(self, mock_time):
 585         mock_time.time.side_effect = count()
 586
 587         auto_connect_ssh = AutoConnectSSH('user1', 'host1', wait=10)
 588         auto_connect_ssh._get_client = mock__get_client = mock.Mock()
 589         mock__get_client.side_effect = exceptions.SSHError
 590
 591         with self.assertRaises(exceptions.SSHTimeout):
 592             auto_connect_ssh._connect()
 593
 594         self.assertEqual(mock_time.time.call_count, 12)
 595
 596     def test_get_file_obj(self):
 597         auto_connect_ssh = AutoConnectSSH('user1', 'host1', wait=10)
 598         auto_connect_ssh._get_client = mock__get_client = mock.Mock()
 599         mock_client = mock__get_client()
 600         mock_open_sftp = mock_client.open_sftp()
 601         mock_sftp = mock.Mock()
 602         mock_open_sftp.__enter__ = mock.Mock(return_value=mock_sftp)
 603         mock_open_sftp.__exit__ = mock.Mock()
 604
 605         auto_connect_ssh.get_file_obj('remote/path', mock.Mock())
 606
 607         mock_sftp.getfo.assert_called_once()
 608
 609     def test__make_dict(self):
 610         auto_connect_ssh = AutoConnectSSH('user1', 'host1')
 611
 612         expected = {
 613             'user': 'user1',
 614             'host': 'host1',
 615             'port': SSH.SSH_PORT,
 616             'pkey': None,
 617             'key_filename': None,
 618             'password': None,
 619             'name': None,
 620             'wait': AutoConnectSSH.DEFAULT_WAIT_TIMEOUT,
 621         }
 622         result = auto_connect_ssh._make_dict()
 623         self.assertDictEqual(result, expected)
 624
 625     def test_get_class(self):
 626         auto_connect_ssh = AutoConnectSSH('user1', 'host1')
 627
 628         self.assertEqual(auto_connect_ssh.get_class(), AutoConnectSSH)
 629
 630     def test_drop_connection(self):
 631         auto_connect_ssh = AutoConnectSSH('user1', 'host1')
 632         self.assertFalse(auto_connect_ssh._client)
 633         auto_connect_ssh._client = True
 634         auto_connect_ssh.drop_connection()
 635         self.assertFalse(auto_connect_ssh._client)
 636
 637     @mock.patch('yardstick.ssh.SCPClient')
 638     def test_put(self, mock_scp_client_type):
 639         auto_connect_ssh = AutoConnectSSH('user1', 'host1')
 640         auto_connect_ssh._client = mock.Mock()
 641
 642         auto_connect_ssh.put('a', 'z')
 643         with mock_scp_client_type() as mock_scp_client:
 644             mock_scp_client.put.assert_called_once()
 645
 646     @mock.patch('yardstick.ssh.SCPClient')
 647     def test_get(self, mock_scp_client_type):
 648         auto_connect_ssh = AutoConnectSSH('user1', 'host1')
 649         auto_connect_ssh._client = mock.Mock()
 650
 651         auto_connect_ssh.get('a', 'z')
 652         with mock_scp_client_type() as mock_scp_client:
 653             mock_scp_client.get.assert_called_once()
 654
 655     def test_put_file(self):
 656         auto_connect_ssh = AutoConnectSSH('user1', 'host1')
 657         auto_connect_ssh._client = mock.Mock()
 658         auto_connect_ssh._put_file_sftp = mock_put_sftp = mock.Mock()
 659
 660         auto_connect_ssh.put_file('a', 'b')
 661         mock_put_sftp.assert_called_once()
 662
 663     def test_execute(self):
 664         auto_connect_ssh = AutoConnectSSH('user1', 'host1')
 665         auto_connect_ssh._client = mock.Mock()
 666         auto_connect_ssh.run = mock.Mock(return_value=0)
 667         exit_code, _, _ = auto_connect_ssh.execute('')
 668         self.assertEqual(exit_code, 0)
 669
 670     def _mock_run(self, *args, **kwargs):
 671         if args[0] == 'ls':
 672             if kwargs.get('raise_on_error'):
 673                 raise exceptions.SSHError(error_msg='Command error')
 674             return 1
 675         return 0
 676
 677     def test_execute_command_error(self):
 678         auto_connect_ssh = AutoConnectSSH('user1', 'host1')
 679         auto_connect_ssh._client = mock.Mock()
 680         auto_connect_ssh.run = mock.Mock(side_effect=self._mock_run)
 681         self.assertRaises(exceptions.SSHError, auto_connect_ssh.execute, 'ls',
 682                           raise_on_error=True)
 683         exit_code, _, _ = auto_connect_ssh.execute('ls')
 684         self.assertNotEqual(exit_code, 0)