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