blob: cdb18efb5b512909a4715f651d232416f59279ef [file] [log] [blame]
Fang Deng69498c32017-03-02 14:29:30 -08001#!/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 Deng69498c32017-03-02 14:29:30 -080016"""Tests for acloud.internal.lib.utils."""
17
Fang Deng26e4dc12018-03-04 19:01:59 -080018import errno
Fang Deng69498c32017-03-02 14:29:30 -080019import getpass
20import os
Fang Deng26e4dc12018-03-04 19:01:59 -080021import shutil
Fang Deng69498c32017-03-02 14:29:30 -080022import subprocess
Fang Deng26e4dc12018-03-04 19:01:59 -080023import tempfile
Fang Dengf24be082018-02-10 10:09:55 -080024import time
Sam Chiu7a477f52018-10-22 11:20:36 +080025import Tkinter
Fang Deng69498c32017-03-02 14:29:30 -080026
cylan0d77ae12018-05-18 08:36:48 +000027import unittest
Fang Deng69498c32017-03-02 14:29:30 -080028import mock
29
Fang Deng69498c32017-03-02 14:29:30 -080030from acloud.internal.lib import driver_test_lib
31from acloud.internal.lib import utils
32
33
34class UtilsTest(driver_test_lib.BaseDriverTest):
cylan0d77ae12018-05-18 08:36:48 +000035 """Test Utils."""
Fang Deng69498c32017-03-02 14:29:30 -080036
cylan0d77ae12018-05-18 08:36:48 +000037 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 Deng26e4dc12018-03-04 19:01:59 -080047
cylan0d77ae12018-05-18 08:36:48 +000048 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 Deng26e4dc12018-03-04 19:01:59 -080053
cylan0d77ae12018-05-18 08:36:48 +000054 class ExpectedException(Exception):
55 """Expected exception."""
56 pass
Fang Deng26e4dc12018-03-04 19:01:59 -080057
cylan0d77ae12018-05-18 08:36:48 +000058 def _Call():
59 with utils.TempDir():
60 raise ExpectedException("Expected exception.")
Fang Deng26e4dc12018-03-04 19:01:59 -080061
cylan0d77ae12018-05-18 08:36:48 +000062 # 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 Deng26e4dc12018-03-04 19:01:59 -080066
cylan0d77ae12018-05-18 08:36:48 +000067 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 Deng26e4dc12018-03-04 19:01:59 -080074
cylan0d77ae12018-05-18 08:36:48 +000075 def _Call():
76 with utils.TempDir():
77 pass
Fang Deng26e4dc12018-03-04 19:01:59 -080078
cylan0d77ae12018-05-18 08:36:48 +000079 # 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 Deng26e4dc12018-03-04 19:01:59 -080085
cylan0d77ae12018-05-18 08:36:48 +000086 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 Deng26e4dc12018-03-04 19:01:59 -080092
cylan0d77ae12018-05-18 08:36:48 +000093 def _Call():
94 with utils.TempDir():
95 pass
Fang Deng26e4dc12018-03-04 19:01:59 -080096
cylan0d77ae12018-05-18 08:36:48 +000097 # 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 Deng26e4dc12018-03-04 19:01:59 -0800101
cylan0d77ae12018-05-18 08:36:48 +0000102 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 Deng69498c32017-03-02 14:29:30 -0800108
cylan0d77ae12018-05-18 08:36:48 +0000109 class ExpectedException(Exception):
110 """Expected exception."""
111 pass
Fang Deng69498c32017-03-02 14:29:30 -0800112
cylan0d77ae12018-05-18 08:36:48 +0000113 def _Call():
114 with utils.TempDir():
115 raise ExpectedException("Expected Exception")
Fang Dengf24be082018-02-10 10:09:55 -0800116
cylan0d77ae12018-05-18 08:36:48 +0000117 # 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 Dengf24be082018-02-10 10:09:55 -0800123
cylan0d77ae12018-05-18 08:36:48 +0000124 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"
cylan4f73c1f2018-07-19 16:40:31 +0800128 self.Patch(os.path, "exists", side_effect=[True, True])
cylan0d77ae12018-05-18 08:36:48 +0000129 self.Patch(subprocess, "check_call")
cylan4f73c1f2018-07-19 16:40:31 +0800130 self.Patch(os, "makedirs", return_value=True)
cylan0d77ae12018-05-18 08:36:48 +0000131 utils.CreateSshKeyPairIfNotExist(private_key, public_key)
132 self.assertEqual(subprocess.check_call.call_count, 0) #pylint: disable=no-member
Fang Dengf24be082018-02-10 10:09:55 -0800133
cylan0d77ae12018-05-18 08:36:48 +0000134 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)
cylan4f73c1f2018-07-19 16:40:31 +0800139 self.Patch(os, "makedirs", return_value=True)
cylan0d77ae12018-05-18 08:36:48 +0000140 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 Dengf24be082018-02-10 10:09:55 -0800149
cylan4f73c1f2018-07-19 16:40:31 +0800150 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)
cylan4f73c1f2018-07-19 16:40:31 +0800157 self.Patch(subprocess, "check_output")
158 self.Patch(os, "rename")
Kevin Chengda4f07a2018-06-26 10:25:05 -0700159 with mock.patch("__builtin__.open", mock_open):
160 utils.CreateSshKeyPairIfNotExist(private_key, public_key)
cylan4f73c1f2018-07-19 16:40:31 +0800161 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
cylan0d77ae12018-05-18 08:36:48 +0000165 def TestRetryOnException(self):
166 """Test Retry."""
Fang Dengf24be082018-02-10 10:09:55 -0800167
cylan0d77ae12018-05-18 08:36:48 +0000168 def _IsValueError(exc):
169 return isinstance(exc, ValueError)
Fang Dengf24be082018-02-10 10:09:55 -0800170
cylan0d77ae12018-05-18 08:36:48 +0000171 num_retry = 5
Fang Dengf24be082018-02-10 10:09:55 -0800172
cylan0d77ae12018-05-18 08:36:48 +0000173 @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 Chengd25feee2018-05-24 10:15:20 -0700196 0, # sleep_multiplier
197 1, # retry_backoff_factor
cylan0d77ae12018-05-18 08:36:48 +0000198 sentinel=sentinel)
199 self.assertEqual(1 + num_retry, sentinel.alert.call_count)
200
201 def testRetry(self):
202 """Test Retry."""
Kevin Chengd25feee2018-05-24 10:15:20 -0700203 mock_sleep = self.Patch(time, "sleep")
cylan0d77ae12018-05-18 08:36:48 +0000204
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 Chengd25feee2018-05-24 10:15:20 -0700216 1, # sleep_multiplier
217 2, # retry_backoff_factor
cylan0d77ae12018-05-18 08:36:48 +0000218 sentinel=sentinel)
219
220 self.assertEqual(1 + num_retry, sentinel.alert.call_count)
Kevin Chengd25feee2018-05-24 10:15:20 -0700221 mock_sleep.assert_has_calls(
cylan0d77ae12018-05-18 08:36:48 +0000222 [
223 mock.call(1),
224 mock.call(2),
225 mock.call(4),
226 mock.call(8),
227 mock.call(16)
228 ])
Fang Dengf24be082018-02-10 10:09:55 -0800229
Kevin Chengeb85e862018-10-09 15:35:13 -0700230 @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 Chiu7a477f52018-10-22 11:20:36 +0800248 @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 Deng69498c32017-03-02 14:29:30 -0800282
283if __name__ == "__main__":
284 unittest.main()