Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 1 | #!/usr/bin/env python |
| 2 | # |
| 3 | # Copyright 2016 - The Android Open Source Project |
| 4 | # |
| 5 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 6 | # you may not use this file except in compliance with the License. |
| 7 | # You may obtain a copy of the License at |
| 8 | # |
| 9 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | # |
| 11 | # Unless required by applicable law or agreed to in writing, software |
| 12 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | # See the License for the specific language governing permissions and |
| 15 | # limitations under the License. |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 16 | """Tests for acloud.internal.lib.utils.""" |
| 17 | |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 18 | import errno |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 19 | import getpass |
| 20 | import os |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 21 | import shutil |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 22 | import subprocess |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 23 | import tempfile |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 24 | import time |
Sam Chiu | 7a477f5 | 2018-10-22 11:20:36 +0800 | [diff] [blame^] | 25 | import Tkinter |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 26 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 27 | import unittest |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 28 | import mock |
| 29 | |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 30 | from acloud.internal.lib import driver_test_lib |
| 31 | from acloud.internal.lib import utils |
| 32 | |
| 33 | |
| 34 | class UtilsTest(driver_test_lib.BaseDriverTest): |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 35 | """Test Utils.""" |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 36 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 37 | def TestTempDirSuccess(self): |
| 38 | """Test create a temp dir.""" |
| 39 | self.Patch(os, "chmod") |
| 40 | self.Patch(tempfile, "mkdtemp", return_value="/tmp/tempdir") |
| 41 | self.Patch(shutil, "rmtree") |
| 42 | with utils.TempDir(): |
| 43 | pass |
| 44 | # Verify. |
| 45 | tempfile.mkdtemp.assert_called_once() # pylint: disable=no-member |
| 46 | shutil.rmtree.assert_called_with("/tmp/tempdir") # pylint: disable=no-member |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 47 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 48 | def TestTempDirExceptionRaised(self): |
| 49 | """Test create a temp dir and exception is raised within with-clause.""" |
| 50 | self.Patch(os, "chmod") |
| 51 | self.Patch(tempfile, "mkdtemp", return_value="/tmp/tempdir") |
| 52 | self.Patch(shutil, "rmtree") |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 53 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 54 | class ExpectedException(Exception): |
| 55 | """Expected exception.""" |
| 56 | pass |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 57 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 58 | def _Call(): |
| 59 | with utils.TempDir(): |
| 60 | raise ExpectedException("Expected exception.") |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 61 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 62 | # Verify. ExpectedException should be raised. |
| 63 | self.assertRaises(ExpectedException, _Call) |
| 64 | tempfile.mkdtemp.assert_called_once() # pylint: disable=no-member |
| 65 | shutil.rmtree.assert_called_with("/tmp/tempdir") #pylint: disable=no-member |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 66 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 67 | def testTempDirWhenDeleteTempDirNoLongerExist(self): # pylint: disable=invalid-name |
| 68 | """Test create a temp dir and dir no longer exists during deletion.""" |
| 69 | self.Patch(os, "chmod") |
| 70 | self.Patch(tempfile, "mkdtemp", return_value="/tmp/tempdir") |
| 71 | expected_error = EnvironmentError() |
| 72 | expected_error.errno = errno.ENOENT |
| 73 | self.Patch(shutil, "rmtree", side_effect=expected_error) |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 74 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 75 | def _Call(): |
| 76 | with utils.TempDir(): |
| 77 | pass |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 78 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 79 | # Verify no exception should be raised when rmtree raises |
| 80 | # EnvironmentError with errno.ENOENT, i.e. |
| 81 | # directory no longer exists. |
| 82 | _Call() |
| 83 | tempfile.mkdtemp.assert_called_once() #pylint: disable=no-member |
| 84 | shutil.rmtree.assert_called_with("/tmp/tempdir") #pylint: disable=no-member |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 85 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 86 | def testTempDirWhenDeleteEncounterError(self): |
| 87 | """Test create a temp dir and encoutered error during deletion.""" |
| 88 | self.Patch(os, "chmod") |
| 89 | self.Patch(tempfile, "mkdtemp", return_value="/tmp/tempdir") |
| 90 | expected_error = OSError("Expected OS Error") |
| 91 | self.Patch(shutil, "rmtree", side_effect=expected_error) |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 92 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 93 | def _Call(): |
| 94 | with utils.TempDir(): |
| 95 | pass |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 96 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 97 | # Verify OSError should be raised. |
| 98 | self.assertRaises(OSError, _Call) |
| 99 | tempfile.mkdtemp.assert_called_once() #pylint: disable=no-member |
| 100 | shutil.rmtree.assert_called_with("/tmp/tempdir") #pylint: disable=no-member |
Fang Deng | 26e4dc1 | 2018-03-04 19:01:59 -0800 | [diff] [blame] | 101 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 102 | def testTempDirOrininalErrorRaised(self): |
| 103 | """Test original error is raised even if tmp dir deletion failed.""" |
| 104 | self.Patch(os, "chmod") |
| 105 | self.Patch(tempfile, "mkdtemp", return_value="/tmp/tempdir") |
| 106 | expected_error = OSError("Expected OS Error") |
| 107 | self.Patch(shutil, "rmtree", side_effect=expected_error) |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 108 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 109 | class ExpectedException(Exception): |
| 110 | """Expected exception.""" |
| 111 | pass |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 112 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 113 | def _Call(): |
| 114 | with utils.TempDir(): |
| 115 | raise ExpectedException("Expected Exception") |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 116 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 117 | # Verify. |
| 118 | # ExpectedException should be raised, and OSError |
| 119 | # should not be raised. |
| 120 | self.assertRaises(ExpectedException, _Call) |
| 121 | tempfile.mkdtemp.assert_called_once() #pylint: disable=no-member |
| 122 | shutil.rmtree.assert_called_with("/tmp/tempdir") #pylint: disable=no-member |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 123 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 124 | def testCreateSshKeyPairKeyAlreadyExists(self): #pylint: disable=invalid-name |
| 125 | """Test when the key pair already exists.""" |
| 126 | public_key = "/fake/public_key" |
| 127 | private_key = "/fake/private_key" |
cylan | 4f73c1f | 2018-07-19 16:40:31 +0800 | [diff] [blame] | 128 | self.Patch(os.path, "exists", side_effect=[True, True]) |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 129 | self.Patch(subprocess, "check_call") |
cylan | 4f73c1f | 2018-07-19 16:40:31 +0800 | [diff] [blame] | 130 | self.Patch(os, "makedirs", return_value=True) |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 131 | utils.CreateSshKeyPairIfNotExist(private_key, public_key) |
| 132 | self.assertEqual(subprocess.check_call.call_count, 0) #pylint: disable=no-member |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 133 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 134 | def testCreateSshKeyPairKeyAreCreated(self): |
| 135 | """Test when the key pair created.""" |
| 136 | public_key = "/fake/public_key" |
| 137 | private_key = "/fake/private_key" |
| 138 | self.Patch(os.path, "exists", return_value=False) |
cylan | 4f73c1f | 2018-07-19 16:40:31 +0800 | [diff] [blame] | 139 | self.Patch(os, "makedirs", return_value=True) |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 140 | self.Patch(subprocess, "check_call") |
| 141 | self.Patch(os, "rename") |
| 142 | utils.CreateSshKeyPairIfNotExist(private_key, public_key) |
| 143 | self.assertEqual(subprocess.check_call.call_count, 1) #pylint: disable=no-member |
| 144 | subprocess.check_call.assert_called_with( #pylint: disable=no-member |
| 145 | utils.SSH_KEYGEN_CMD + |
| 146 | ["-C", getpass.getuser(), "-f", private_key], |
| 147 | stdout=mock.ANY, |
| 148 | stderr=mock.ANY) |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 149 | |
cylan | 4f73c1f | 2018-07-19 16:40:31 +0800 | [diff] [blame] | 150 | def testCreatePublicKeyAreCreated(self): |
| 151 | """Test when the PublicKey created.""" |
| 152 | public_key = "/fake/public_key" |
| 153 | private_key = "/fake/private_key" |
| 154 | self.Patch(os.path, "exists", side_effect=[False, True, True]) |
| 155 | self.Patch(os, "makedirs", return_value=True) |
| 156 | mock_open = mock.mock_open(read_data=public_key) |
cylan | 4f73c1f | 2018-07-19 16:40:31 +0800 | [diff] [blame] | 157 | self.Patch(subprocess, "check_output") |
| 158 | self.Patch(os, "rename") |
Kevin Cheng | da4f07a | 2018-06-26 10:25:05 -0700 | [diff] [blame] | 159 | with mock.patch("__builtin__.open", mock_open): |
| 160 | utils.CreateSshKeyPairIfNotExist(private_key, public_key) |
cylan | 4f73c1f | 2018-07-19 16:40:31 +0800 | [diff] [blame] | 161 | self.assertEqual(subprocess.check_output.call_count, 1) #pylint: disable=no-member |
| 162 | subprocess.check_output.assert_called_with( #pylint: disable=no-member |
| 163 | utils.SSH_KEYGEN_PUB_CMD +["-f", private_key]) |
| 164 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 165 | def TestRetryOnException(self): |
| 166 | """Test Retry.""" |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 167 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 168 | def _IsValueError(exc): |
| 169 | return isinstance(exc, ValueError) |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 170 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 171 | num_retry = 5 |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 172 | |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 173 | @utils.RetryOnException(_IsValueError, num_retry) |
| 174 | def _RaiseAndRetry(sentinel): |
| 175 | sentinel.alert() |
| 176 | raise ValueError("Fake error.") |
| 177 | |
| 178 | sentinel = mock.MagicMock() |
| 179 | self.assertRaises(ValueError, _RaiseAndRetry, sentinel) |
| 180 | self.assertEqual(1 + num_retry, sentinel.alert.call_count) |
| 181 | |
| 182 | def testRetryExceptionType(self): |
| 183 | """Test RetryExceptionType function.""" |
| 184 | |
| 185 | def _RaiseAndRetry(sentinel): |
| 186 | sentinel.alert() |
| 187 | raise ValueError("Fake error.") |
| 188 | |
| 189 | num_retry = 5 |
| 190 | sentinel = mock.MagicMock() |
| 191 | self.assertRaises( |
| 192 | ValueError, |
| 193 | utils.RetryExceptionType, (KeyError, ValueError), |
| 194 | num_retry, |
| 195 | _RaiseAndRetry, |
Kevin Cheng | d25feee | 2018-05-24 10:15:20 -0700 | [diff] [blame] | 196 | 0, # sleep_multiplier |
| 197 | 1, # retry_backoff_factor |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 198 | sentinel=sentinel) |
| 199 | self.assertEqual(1 + num_retry, sentinel.alert.call_count) |
| 200 | |
| 201 | def testRetry(self): |
| 202 | """Test Retry.""" |
Kevin Cheng | d25feee | 2018-05-24 10:15:20 -0700 | [diff] [blame] | 203 | mock_sleep = self.Patch(time, "sleep") |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 204 | |
| 205 | def _RaiseAndRetry(sentinel): |
| 206 | sentinel.alert() |
| 207 | raise ValueError("Fake error.") |
| 208 | |
| 209 | num_retry = 5 |
| 210 | sentinel = mock.MagicMock() |
| 211 | self.assertRaises( |
| 212 | ValueError, |
| 213 | utils.RetryExceptionType, (ValueError, KeyError), |
| 214 | num_retry, |
| 215 | _RaiseAndRetry, |
Kevin Cheng | d25feee | 2018-05-24 10:15:20 -0700 | [diff] [blame] | 216 | 1, # sleep_multiplier |
| 217 | 2, # retry_backoff_factor |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 218 | sentinel=sentinel) |
| 219 | |
| 220 | self.assertEqual(1 + num_retry, sentinel.alert.call_count) |
Kevin Cheng | d25feee | 2018-05-24 10:15:20 -0700 | [diff] [blame] | 221 | mock_sleep.assert_has_calls( |
cylan | 0d77ae1 | 2018-05-18 08:36:48 +0000 | [diff] [blame] | 222 | [ |
| 223 | mock.call(1), |
| 224 | mock.call(2), |
| 225 | mock.call(4), |
| 226 | mock.call(8), |
| 227 | mock.call(16) |
| 228 | ]) |
Fang Deng | f24be08 | 2018-02-10 10:09:55 -0800 | [diff] [blame] | 229 | |
Kevin Cheng | eb85e86 | 2018-10-09 15:35:13 -0700 | [diff] [blame] | 230 | @mock.patch("__builtin__.raw_input") |
| 231 | def testGetAnswerFromList(self, mock_raw_input): |
| 232 | """Test GetAnswerFromList.""" |
| 233 | answer_list = ["image1.zip", "image2.zip", "image3.zip"] |
| 234 | mock_raw_input.return_value = 0 |
| 235 | with self.assertRaises(SystemExit): |
| 236 | utils.GetAnswerFromList(answer_list) |
| 237 | mock_raw_input.side_effect = [1, 2, 3, 1] |
| 238 | self.assertEqual(utils.GetAnswerFromList(answer_list), |
| 239 | ["image1.zip"]) |
| 240 | self.assertEqual(utils.GetAnswerFromList(answer_list), |
| 241 | ["image2.zip"]) |
| 242 | self.assertEqual(utils.GetAnswerFromList(answer_list), |
| 243 | ["image3.zip"]) |
| 244 | self.assertEqual(utils.GetAnswerFromList(answer_list, |
| 245 | enable_choose_all=True), |
| 246 | answer_list) |
| 247 | |
Sam Chiu | 7a477f5 | 2018-10-22 11:20:36 +0800 | [diff] [blame^] | 248 | @mock.patch.object(Tkinter.Tk, "winfo_screenwidth") |
| 249 | @mock.patch.object(Tkinter.Tk, "winfo_screenheight") |
| 250 | def testCalculateVNCScreenRatio(self, mock_screenheight, mock_screenwidth): |
| 251 | """Test Calculating the scale ratio of VNC display.""" |
| 252 | |
| 253 | # Get scale-down ratio if screen height is smaller than AVD height. |
| 254 | mock_screenheight.return_value = 800 |
| 255 | mock_screenwidth.return_value = 1200 |
| 256 | avd_h = 1920 |
| 257 | avd_w = 1080 |
| 258 | self.assertEqual(utils.CalculateVNCScreenRatio(avd_w, avd_h), 0.4) |
| 259 | |
| 260 | # Get scale-down ratio if screen width is smaller than AVD width. |
| 261 | mock_screenheight.return_value = 800 |
| 262 | mock_screenwidth.return_value = 1200 |
| 263 | avd_h = 900 |
| 264 | avd_w = 1920 |
| 265 | self.assertEqual(utils.CalculateVNCScreenRatio(avd_w, avd_h), 0.6) |
| 266 | |
| 267 | # Scale ratio = 1 if screen is larger than AVD. |
| 268 | mock_screenheight.return_value = 1080 |
| 269 | mock_screenwidth.return_value = 1920 |
| 270 | avd_h = 800 |
| 271 | avd_w = 1280 |
| 272 | self.assertEqual(utils.CalculateVNCScreenRatio(avd_w, avd_h), 1) |
| 273 | |
| 274 | # Get the scale if ratio of width is smaller than the |
| 275 | # ratio of height. |
| 276 | mock_screenheight.return_value = 1200 |
| 277 | mock_screenwidth.return_value = 800 |
| 278 | avd_h = 1920 |
| 279 | avd_w = 1080 |
| 280 | self.assertEqual(utils.CalculateVNCScreenRatio(avd_w, avd_h), 0.6) |
| 281 | |
Fang Deng | 69498c3 | 2017-03-02 14:29:30 -0800 | [diff] [blame] | 282 | |
| 283 | if __name__ == "__main__": |
| 284 | unittest.main() |