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