1 #!/usr/bin/env python2 2 # 3 # Copyright 2014 Google Inc. All Rights Reserved 4 """Download image unittest.""" 5 6 from __future__ import print_function 7 8 import os 9 import mock 10 import unittest 11 12 import download_images 13 from cros_utils import command_executer 14 from cros_utils import logger 15 16 import test_flag 17 18 MOCK_LOGGER = logger.GetLogger(log_dir='', mock=True) 19 20 21 class ImageDownloaderTestcast(unittest.TestCase): 22 """The image downloader test class.""" 23 24 def __init__(self, *args, **kwargs): 25 super(ImageDownloaderTestcast, self).__init__(*args, **kwargs) 26 self.called_download_image = False 27 self.called_uncompress_image = False 28 self.called_get_build_id = False 29 self.called_download_autotest_files = False 30 31 @mock.patch.object(os, 'makedirs') 32 @mock.patch.object(os.path, 'exists') 33 def test_download_image(self, mock_path_exists, mock_mkdirs): 34 35 # Set mock and test values. 36 mock_cmd_exec = mock.Mock(spec=command_executer.CommandExecuter) 37 test_chroot = '/usr/local/home/chromeos' 38 test_build_id = 'lumpy-release/R36-5814.0.0' 39 image_path = ('gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz' 40 % test_build_id) 41 42 downloader = download_images.ImageDownloader( 43 logger_to_use=MOCK_LOGGER, cmd_exec=mock_cmd_exec) 44 45 # Set os.path.exists to always return False and run downloader 46 mock_path_exists.return_value = False 47 test_flag.SetTestMode(True) 48 self.assertRaises(download_images.MissingImage, downloader.DownloadImage, 49 test_chroot, test_build_id, image_path) 50 51 # Verify os.path.exists was called twice, with proper arguments. 52 self.assertEqual(mock_path_exists.call_count, 2) 53 mock_path_exists.assert_called_with( 54 '/usr/local/home/chromeos/chroot/tmp/lumpy-release/' 55 'R36-5814.0.0/chromiumos_test_image.bin') 56 mock_path_exists.assert_any_call( 57 '/usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0') 58 59 # Verify we called os.mkdirs 60 self.assertEqual(mock_mkdirs.call_count, 1) 61 mock_mkdirs.assert_called_with( 62 '/usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0') 63 64 # Verify we called RunCommand once, with proper arguments. 65 self.assertEqual(mock_cmd_exec.RunCommand.call_count, 1) 66 expected_args = ( 67 '/usr/local/home/chromeos/chromium/tools/depot_tools/gsutil.py ' 68 'cp gs://chromeos-image-archive/lumpy-release/R36-5814.0.0/' 69 'chromiumos_test_image.tar.xz ' 70 '/usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0') 71 72 mock_cmd_exec.RunCommand.assert_called_with(expected_args) 73 74 # Reset the velues in the mocks; set os.path.exists to always return True. 75 mock_path_exists.reset_mock() 76 mock_cmd_exec.reset_mock() 77 mock_path_exists.return_value = True 78 79 # Run downloader 80 downloader.DownloadImage(test_chroot, test_build_id, image_path) 81 82 # Verify os.path.exists was called twice, with proper arguments. 83 self.assertEqual(mock_path_exists.call_count, 2) 84 mock_path_exists.assert_called_with( 85 '/usr/local/home/chromeos/chroot/tmp/lumpy-release/' 86 'R36-5814.0.0/chromiumos_test_image.bin') 87 mock_path_exists.assert_any_call( 88 '/usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0') 89 90 # Verify we made no RunCommand or ChrootRunCommand calls (since 91 # os.path.exists returned True, there was no work do be done). 92 self.assertEqual(mock_cmd_exec.RunCommand.call_count, 0) 93 self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 0) 94 95 @mock.patch.object(os.path, 'exists') 96 def test_uncompress_image(self, mock_path_exists): 97 98 # set mock and test values. 99 mock_cmd_exec = mock.Mock(spec=command_executer.CommandExecuter) 100 test_chroot = '/usr/local/home/chromeos' 101 test_build_id = 'lumpy-release/R36-5814.0.0' 102 103 downloader = download_images.ImageDownloader( 104 logger_to_use=MOCK_LOGGER, cmd_exec=mock_cmd_exec) 105 106 # Set os.path.exists to always return False and run uncompress. 107 mock_path_exists.return_value = False 108 self.assertRaises(download_images.MissingImage, downloader.UncompressImage, 109 test_chroot, test_build_id) 110 111 # Verify os.path.exists was called once, with correct arguments. 112 self.assertEqual(mock_path_exists.call_count, 1) 113 mock_path_exists.assert_called_with( 114 '/usr/local/home/chromeos/chroot/tmp/lumpy-release/' 115 'R36-5814.0.0/chromiumos_test_image.bin') 116 117 # Verify RunCommand was called twice with correct arguments. 118 self.assertEqual(mock_cmd_exec.RunCommand.call_count, 2) 119 # Call 1, should have 2 arguments 120 self.assertEqual(len(mock_cmd_exec.RunCommand.call_args_list[0]), 2) 121 actual_arg = mock_cmd_exec.RunCommand.call_args_list[0][0] 122 expected_arg = ( 123 'cd /usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0 ; ' 124 'tar -Jxf chromiumos_test_image.tar.xz ',) 125 self.assertEqual(expected_arg, actual_arg) 126 # 2nd arg must be exception handler 127 except_handler_string = 'RunCommandExceptionHandler.HandleException' 128 self.assertTrue( 129 except_handler_string in repr( 130 mock_cmd_exec.RunCommand.call_args_list[0][1])) 131 132 # Call 2, should have 2 arguments 133 self.assertEqual(len(mock_cmd_exec.RunCommand.call_args_list[1]), 2) 134 actual_arg = mock_cmd_exec.RunCommand.call_args_list[1][0] 135 expected_arg = ( 136 'cd /usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0 ; ' 137 'rm -f chromiumos_test_image.bin ',) 138 self.assertEqual(expected_arg, actual_arg) 139 # 2nd arg must be empty 140 self.assertTrue('{}' in repr(mock_cmd_exec.RunCommand.call_args_list[1][1])) 141 142 # Set os.path.exists to always return True and run uncompress. 143 mock_path_exists.reset_mock() 144 mock_cmd_exec.reset_mock() 145 mock_path_exists.return_value = True 146 downloader.UncompressImage(test_chroot, test_build_id) 147 148 # Verify os.path.exists was called once, with correct arguments. 149 self.assertEqual(mock_path_exists.call_count, 1) 150 mock_path_exists.assert_called_with( 151 '/usr/local/home/chromeos/chroot/tmp/lumpy-release/' 152 'R36-5814.0.0/chromiumos_test_image.bin') 153 154 # Verify RunCommand was not called. 155 self.assertEqual(mock_cmd_exec.RunCommand.call_count, 0) 156 157 def test_run(self): 158 159 # Set test arguments 160 test_chroot = '/usr/local/home/chromeos' 161 test_build_id = 'remote/lumpy/latest-dev' 162 test_empty_autotest_path = '' 163 test_autotest_path = '/tmp/autotest' 164 165 # Set values to test/check. 166 self.called_download_image = False 167 self.called_uncompress_image = False 168 self.called_get_build_id = False 169 self.called_download_autotest_files = False 170 171 # Define fake stub functions for Run to call 172 def FakeGetBuildID(unused_root, unused_xbuddy_label): 173 self.called_get_build_id = True 174 return 'lumpy-release/R36-5814.0.0' 175 176 def GoodDownloadImage(root, build_id, image_path): 177 if root or build_id or image_path: 178 pass 179 self.called_download_image = True 180 return 'chromiumos_test_image.bin' 181 182 def BadDownloadImage(root, build_id, image_path): 183 if root or build_id or image_path: 184 pass 185 self.called_download_image = True 186 raise download_images.MissingImage('Could not download image') 187 188 def FakeUncompressImage(root, build_id): 189 if root or build_id: 190 pass 191 self.called_uncompress_image = True 192 return 0 193 194 def FakeDownloadAutotestFiles(root, build_id): 195 if root or build_id: 196 pass 197 self.called_download_autotest_files = True 198 return 'autotest' 199 200 # Initialize downloader 201 downloader = download_images.ImageDownloader(logger_to_use=MOCK_LOGGER) 202 203 # Set downloader to call fake stubs. 204 downloader.GetBuildID = FakeGetBuildID 205 downloader.UncompressImage = FakeUncompressImage 206 downloader.DownloadImage = GoodDownloadImage 207 downloader.DownloadAutotestFiles = FakeDownloadAutotestFiles 208 209 # Call Run. 210 image_path, autotest_path = downloader.Run(test_chroot, test_build_id, 211 test_empty_autotest_path) 212 213 # Make sure it called both _DownloadImage and _UncompressImage 214 self.assertTrue(self.called_download_image) 215 self.assertTrue(self.called_uncompress_image) 216 # Make sure it called DownloadAutotestFiles 217 self.assertTrue(self.called_download_autotest_files) 218 # Make sure it returned an image and autotest path returned from this call 219 self.assertTrue(image_path == 'chromiumos_test_image.bin') 220 self.assertTrue(autotest_path == 'autotest') 221 222 # Call Run with a non-empty autotest path 223 self.called_download_autotest_files = False 224 225 image_path, autotest_path = downloader.Run(test_chroot, test_build_id, 226 test_autotest_path) 227 228 # Verify that downloadAutotestFiles was not called 229 self.assertFalse(self.called_download_autotest_files) 230 # Make sure it returned the specified autotest path returned from this call 231 self.assertTrue(autotest_path == test_autotest_path) 232 233 # Reset values; Now use fake stub that simulates DownloadImage failing. 234 self.called_download_image = False 235 self.called_uncompress_image = False 236 self.called_download_autotest_files = False 237 downloader.DownloadImage = BadDownloadImage 238 239 # Call Run again. 240 self.assertRaises(download_images.MissingImage, downloader.Run, test_chroot, 241 test_autotest_path, test_build_id) 242 243 # Verify that UncompressImage and downloadAutotestFiles were not called, 244 # since _DownloadImage "failed" 245 self.assertTrue(self.called_download_image) 246 self.assertFalse(self.called_uncompress_image) 247 self.assertFalse(self.called_download_autotest_files) 248 249 250 if __name__ == '__main__': 251 unittest.main() 252
