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 unittest
  21 import mock
  22
  23 from yardstick import ssh
  24
  25
  26 class FakeParamikoException(Exception):
  27     pass
  28
  29
  30 class SSHTestCase(unittest.TestCase):
  31     """Test all small SSH methods."""
  32
  33     def setUp(self):
  34         super(SSHTestCase, self).setUp()
  35         self.test_client = ssh.SSH("root", "example.net")
  36
  37     @mock.patch("yardstick.ssh.SSH._get_pkey")
  38     def test_construct(self, mock_ssh__get_pkey):
  39         mock_ssh__get_pkey.return_value = "pkey"
  40         test_ssh = ssh.SSH("root", "example.net", port=33, pkey="key",
  41                            key_filename="kf", password="secret")
  42         mock_ssh__get_pkey.assert_called_once_with("key")
  43         self.assertEqual("root", test_ssh.user)
  44         self.assertEqual("example.net", test_ssh.host)
  45         self.assertEqual(33, test_ssh.port)
  46         self.assertEqual("pkey", test_ssh.pkey)
  47         self.assertEqual("kf", test_ssh.key_filename)
  48         self.assertEqual("secret", test_ssh.password)
  49
  50     def test_construct_default(self):
  51         self.assertEqual("root", self.test_client.user)
  52         self.assertEqual("example.net", self.test_client.host)
  53         self.assertEqual(22, self.test_client.port)
  54         self.assertIsNone(self.test_client.pkey)
  55         self.assertIsNone(self.test_client.key_filename)
  56         self.assertIsNone(self.test_client.password)
  57
  58     @mock.patch("yardstick.ssh.paramiko")
  59     def test__get_pkey_invalid(self, mock_paramiko):
  60         mock_paramiko.SSHException = FakeParamikoException
  61         rsa = mock_paramiko.rsakey.RSAKey
  62         dss = mock_paramiko.dsskey.DSSKey
  63         rsa.from_private_key.side_effect = mock_paramiko.SSHException
  64         dss.from_private_key.side_effect = mock_paramiko.SSHException
  65         self.assertRaises(ssh.SSHError, self.test_client._get_pkey, "key")
  66
  67     @mock.patch("yardstick.ssh.six.moves.StringIO")
  68     @mock.patch("yardstick.ssh.paramiko")
  69     def test__get_pkey_dss(self, mock_paramiko, mock_string_io):
  70         mock_paramiko.SSHException = FakeParamikoException
  71         mock_string_io.return_value = "string_key"
  72         mock_paramiko.dsskey.DSSKey.from_private_key.return_value = "dss_key"
  73         rsa = mock_paramiko.rsakey.RSAKey
  74         rsa.from_private_key.side_effect = mock_paramiko.SSHException
  75         key = self.test_client._get_pkey("key")
  76         dss_calls = mock_paramiko.dsskey.DSSKey.from_private_key.mock_calls
  77         self.assertEqual([mock.call("string_key")], dss_calls)
  78         self.assertEqual(key, "dss_key")
  79         mock_string_io.assert_called_once_with("key")
  80
  81     @mock.patch("yardstick.ssh.six.moves.StringIO")
  82     @mock.patch("yardstick.ssh.paramiko")
  83     def test__get_pkey_rsa(self, mock_paramiko, mock_string_io):
  84         mock_paramiko.SSHException = FakeParamikoException
  85         mock_string_io.return_value = "string_key"
  86         mock_paramiko.rsakey.RSAKey.from_private_key.return_value = "rsa_key"
  87         dss = mock_paramiko.dsskey.DSSKey
  88         dss.from_private_key.side_effect = mock_paramiko.SSHException
  89         key = self.test_client._get_pkey("key")
  90         rsa_calls = mock_paramiko.rsakey.RSAKey.from_private_key.mock_calls
  91         self.assertEqual([mock.call("string_key")], rsa_calls)
  92         self.assertEqual(key, "rsa_key")
  93         mock_string_io.assert_called_once_with("key")
  94
  95     @mock.patch("yardstick.ssh.SSH._get_pkey")
  96     @mock.patch("yardstick.ssh.paramiko")
  97     def test__get_client(self, mock_paramiko, mock_ssh__get_pkey):
  98         mock_ssh__get_pkey.return_value = "key"
  99         fake_client = mock.Mock()
 100         mock_paramiko.SSHClient.return_value = fake_client
 101         mock_paramiko.AutoAddPolicy.return_value = "autoadd"
 102
 103         test_ssh = ssh.SSH("admin", "example.net", pkey="key")
 104         client = test_ssh._get_client()
 105
 106         self.assertEqual(fake_client, client)
 107         client_calls = [
 108             mock.call.set_missing_host_key_policy("autoadd"),
 109             mock.call.connect("example.net", username="admin",
 110                               port=22, pkey="key", key_filename=None,
 111                               password=None, timeout=1),
 112         ]
 113         self.assertEqual(client_calls, client.mock_calls)
 114
 115     def test_close(self):
 116         with mock.patch.object(self.test_client, "_client") as m_client:
 117             self.test_client.close()
 118         m_client.close.assert_called_once_with()
 119         self.assertFalse(self.test_client._client)
 120
 121     @mock.patch("yardstick.ssh.six.moves.StringIO")
 122     def test_execute(self, mock_string_io):
 123         mock_string_io.side_effect = stdio = [mock.Mock(), mock.Mock()]
 124         stdio[0].read.return_value = "stdout fake data"
 125         stdio[1].read.return_value = "stderr fake data"
 126         with mock.patch.object(self.test_client, "run", return_value=0)\
 127                 as mock_run:
 128             status, stdout, stderr = self.test_client.execute(
 129                 "cmd",
 130                 stdin="fake_stdin",
 131                 timeout=43)
 132         mock_run.assert_called_once_with(
 133             "cmd", stdin="fake_stdin", stdout=stdio[0],
 134             stderr=stdio[1], timeout=43, raise_on_error=False)
 135         self.assertEqual(0, status)
 136         self.assertEqual("stdout fake data", stdout)
 137         self.assertEqual("stderr fake data", stderr)
 138
 139     @mock.patch("yardstick.ssh.time")
 140     def test_wait_timeout(self, mock_time):
 141         mock_time.time.side_effect = [1, 50, 150]
 142         self.test_client.execute = mock.Mock(side_effect=[ssh.SSHError,
 143                                                           ssh.SSHError,
 144                                                           0])
 145         self.assertRaises(ssh.SSHTimeout, self.test_client.wait)
 146         self.assertEqual([mock.call("uname")] * 2,
 147                          self.test_client.execute.mock_calls)
 148
 149     @mock.patch("yardstick.ssh.time")
 150     def test_wait(self, mock_time):
 151         mock_time.time.side_effect = [1, 50, 100]
 152         self.test_client.execute = mock.Mock(side_effect=[ssh.SSHError,
 153                                                           ssh.SSHError,
 154                                                           0])
 155         self.test_client.wait()
 156         self.assertEqual([mock.call("uname")] * 3,
 157                          self.test_client.execute.mock_calls)
 158
 159     @mock.patch("yardstick.ssh.paramiko")
 160     def test_send_command(self, mock_paramiko):
 161         paramiko_sshclient = self.test_client._get_client()
 162         with mock.patch.object(paramiko_sshclient, "exec_command") \
 163             as mock_paramiko_exec_command:
 164             self.test_client.send_command('cmd')
 165         mock_paramiko_exec_command.assert_called_once_with('cmd',
 166                                                             get_pty=True)
 167
 168
 169 class SSHRunTestCase(unittest.TestCase):
 170     """Test SSH.run method in different aspects.
 171
 172     Also tested method "execute".
 173     """
 174
 175     def setUp(self):
 176         super(SSHRunTestCase, self).setUp()
 177
 178         self.fake_client = mock.Mock()
 179         self.fake_session = mock.Mock()
 180         self.fake_transport = mock.Mock()
 181
 182         self.fake_transport.open_session.return_value = self.fake_session
 183         self.fake_client.get_transport.return_value = self.fake_transport
 184
 185         self.fake_session.recv_ready.return_value = False
 186         self.fake_session.recv_stderr_ready.return_value = False
 187         self.fake_session.send_ready.return_value = False
 188         self.fake_session.exit_status_ready.return_value = True
 189         self.fake_session.recv_exit_status.return_value = 0
 190
 191         self.test_client = ssh.SSH("admin", "example.net")
 192         self.test_client._get_client = mock.Mock(return_value=self.fake_client)
 193
 194     @mock.patch("yardstick.ssh.select")
 195     def test_execute(self, mock_select):
 196         mock_select.select.return_value = ([], [], [])
 197         self.fake_session.recv_ready.side_effect = [1, 0, 0]
 198         self.fake_session.recv_stderr_ready.side_effect = [1, 0]
 199         self.fake_session.recv.return_value = "ok"
 200         self.fake_session.recv_stderr.return_value = "error"
 201         self.fake_session.exit_status_ready.return_value = 1
 202         self.fake_session.recv_exit_status.return_value = 127
 203         self.assertEqual((127, "ok", "error"), self.test_client.execute("cmd"))
 204         self.fake_session.exec_command.assert_called_once_with("cmd")
 205
 206     @mock.patch("yardstick.ssh.select")
 207     def test_execute_args(self, mock_select):
 208         mock_select.select.return_value = ([], [], [])
 209         self.fake_session.recv_ready.side_effect = [1, 0, 0]
 210         self.fake_session.recv_stderr_ready.side_effect = [1, 0]
 211         self.fake_session.recv.return_value = "ok"
 212         self.fake_session.recv_stderr.return_value = "error"
 213         self.fake_session.exit_status_ready.return_value = 1
 214         self.fake_session.recv_exit_status.return_value = 127
 215
 216         result = self.test_client.execute("cmd arg1 'arg2 with space'")
 217         self.assertEqual((127, "ok", "error"), result)
 218         self.fake_session.exec_command.assert_called_once_with(
 219             "cmd arg1 'arg2 with space'")
 220
 221     @mock.patch("yardstick.ssh.select")
 222     def test_run(self, mock_select):
 223         mock_select.select.return_value = ([], [], [])
 224         self.assertEqual(0, self.test_client.run("cmd"))
 225
 226     @mock.patch("yardstick.ssh.select")
 227     def test_run_nonzero_status(self, mock_select):
 228         mock_select.select.return_value = ([], [], [])
 229         self.fake_session.recv_exit_status.return_value = 1
 230         self.assertRaises(ssh.SSHError, self.test_client.run, "cmd")
 231         self.assertEqual(1, self.test_client.run("cmd", raise_on_error=False))
 232
 233     @mock.patch("yardstick.ssh.select")
 234     def test_run_stdout(self, mock_select):
 235         mock_select.select.return_value = ([], [], [])
 236         self.fake_session.recv_ready.side_effect = [True, True, False]
 237         self.fake_session.recv.side_effect = ["ok1", "ok2"]
 238         stdout = mock.Mock()
 239         self.test_client.run("cmd", stdout=stdout)
 240         self.assertEqual([mock.call("ok1"), mock.call("ok2")],
 241                          stdout.write.mock_calls)
 242
 243     @mock.patch("yardstick.ssh.select")
 244     def test_run_stderr(self, mock_select):
 245         mock_select.select.return_value = ([], [], [])
 246         self.fake_session.recv_stderr_ready.side_effect = [True, False]
 247         self.fake_session.recv_stderr.return_value = "error"
 248         stderr = mock.Mock()
 249         self.test_client.run("cmd", stderr=stderr)
 250         stderr.write.assert_called_once_with("error")
 251
 252     @mock.patch("yardstick.ssh.select")
 253     def test_run_stdin(self, mock_select):
 254         """Test run method with stdin.
 255
 256         Third send call was called with "e2" because only 3 bytes was sent
 257         by second call. So remainig 2 bytes of "line2" was sent by third call.
 258         """
 259         mock_select.select.return_value = ([], [], [])
 260         self.fake_session.exit_status_ready.side_effect = [0, 0, 0, True]
 261         self.fake_session.send_ready.return_value = True
 262         self.fake_session.send.side_effect = [5, 3, 2]
 263         fake_stdin = mock.Mock()
 264         fake_stdin.read.side_effect = ["line1", "line2", ""]
 265         fake_stdin.closed = False
 266
 267         def close():
 268             fake_stdin.closed = True
 269         fake_stdin.close = mock.Mock(side_effect=close)
 270         self.test_client.run("cmd", stdin=fake_stdin)
 271         call = mock.call
 272         send_calls = [call("line1"), call("line2"), call("e2")]
 273         self.assertEqual(send_calls, self.fake_session.send.mock_calls)
 274
 275     @mock.patch("yardstick.ssh.select")
 276     def test_run_select_error(self, mock_select):
 277         self.fake_session.exit_status_ready.return_value = False
 278         mock_select.select.return_value = ([], [], [True])
 279         self.assertRaises(ssh.SSHError, self.test_client.run, "cmd")
 280
 281     @mock.patch("yardstick.ssh.time")
 282     @mock.patch("yardstick.ssh.select")
 283     def test_run_timemout(self, mock_select, mock_time):
 284         mock_time.time.side_effect = [1, 3700]
 285         mock_select.select.return_value = ([], [], [])
 286         self.fake_session.exit_status_ready.return_value = False
 287         self.assertRaises(ssh.SSHTimeout, self.test_client.run, "cmd")
 288
 289
 290 def main():
 291     unittest.main()
 292
 293 if __name__ == '__main__':
 294     main()