1#!/usr/bin/env python3 2# -*- coding: utf-8 -*- 3# Copyright 2019 The ChromiumOS Authors 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6 7"""Download image unittest.""" 8 9 10import os 11import re 12import unittest 13import unittest.mock as mock 14 15from cros_utils import command_executer 16from cros_utils import logger 17import download_images 18import test_flag 19 20 21MOCK_LOGGER = logger.GetLogger(log_dir="", mock=True) 22 23 24class RegexMatcher: 25 """A regex matcher, for passing to mocks.""" 26 27 def __init__(self, regex): 28 self._regex = re.compile(regex) 29 30 def __eq__(self, string): 31 return self._regex.search(string) is not None 32 33 34class ImageDownloaderTestcast(unittest.TestCase): 35 """The image downloader test class.""" 36 37 def __init__(self, *args, **kwargs): 38 super(ImageDownloaderTestcast, self).__init__(*args, **kwargs) 39 self.called_download_image = False 40 self.called_uncompress_image = False 41 self.called_get_build_id = False 42 self.called_download_autotest_files = False 43 self.called_download_debug_file = False 44 45 @mock.patch.object(os, "makedirs") 46 @mock.patch.object(os.path, "exists") 47 def test_download_image(self, mock_path_exists, mock_mkdirs): 48 # Set mock and test values. 49 mock_cmd_exec = mock.Mock(spec=command_executer.CommandExecuter) 50 test_chroot = "/usr/local/home/chromeos" 51 test_build_id = "lumpy-release/R36-5814.0.0" 52 image_path = ( 53 "gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz" 54 % test_build_id 55 ) 56 57 downloader = download_images.ImageDownloader( 58 logger_to_use=MOCK_LOGGER, cmd_exec=mock_cmd_exec 59 ) 60 61 # Set os.path.exists to always return False and run downloader 62 mock_path_exists.return_value = False 63 test_flag.SetTestMode(True) 64 self.assertRaises( 65 download_images.MissingImage, 66 downloader.DownloadImage, 67 test_chroot, 68 test_build_id, 69 image_path, 70 ) 71 72 # Verify os.path.exists was called thrice, with proper arguments. 73 self.assertEqual(mock_path_exists.call_count, 3) 74 mock_path_exists.assert_any_call( 75 RegexMatcher( 76 "/usr/local/home/chromeos/.*tmp/lumpy-release/" 77 "R36-5814.0.0/chromiumos_test_image.bin" 78 ) 79 ) 80 mock_path_exists.assert_any_call( 81 RegexMatcher( 82 "/usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0" 83 ) 84 ) 85 mock_path_exists.assert_any_call("/etc/cros_chroot_version") 86 87 # Verify we called os.mkdirs 88 self.assertEqual(mock_mkdirs.call_count, 1) 89 mock_mkdirs.assert_called_with( 90 RegexMatcher( 91 "/usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0" 92 ) 93 ) 94 95 # Verify we called RunCommand once, with proper arguments. 96 self.assertEqual(mock_cmd_exec.RunCommand.call_count, 1) 97 expected_args = RegexMatcher( 98 "/usr/local/home/chromeos/src/chromium/depot_tools/gsutil.py " 99 "cp gs://chromeos-image-archive/lumpy-release/R36-5814.0.0/" 100 "chromiumos_test_image.tar.xz " 101 "/usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0" 102 ) 103 104 mock_cmd_exec.RunCommand.assert_called_with(expected_args) 105 106 # Reset the velues in the mocks; set os.path.exists to always return 107 # True (except for "inside chroot" check). 108 mock_path_exists.reset_mock() 109 mock_cmd_exec.reset_mock() 110 mock_path_exists.side_effect = lambda x: x != "/etc/cros_chroot_version" 111 112 # Run downloader 113 downloader.DownloadImage(test_chroot, test_build_id, image_path) 114 115 # Verify os.path.exists was called thrice, with proper arguments. 116 self.assertEqual(mock_path_exists.call_count, 3) 117 mock_path_exists.assert_called_with( 118 RegexMatcher( 119 "/usr/local/home/chromeos/.*tmp/lumpy-release/" 120 "R36-5814.0.0/chromiumos_test_image.bin" 121 ) 122 ) 123 mock_path_exists.assert_any_call( 124 RegexMatcher( 125 "/usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0" 126 ) 127 ) 128 mock_path_exists.assert_any_call("/etc/cros_chroot_version") 129 130 # Verify we made no RunCommand or ChrootRunCommand calls (since 131 # os.path.exists returned True, there was no work do be done). 132 self.assertEqual(mock_cmd_exec.RunCommand.call_count, 0) 133 self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 0) 134 135 @mock.patch.object(os.path, "exists") 136 def test_uncompress_image(self, mock_path_exists): 137 # set mock and test values. 138 mock_cmd_exec = mock.Mock(spec=command_executer.CommandExecuter) 139 test_chroot = "/usr/local/home/chromeos" 140 test_build_id = "lumpy-release/R36-5814.0.0" 141 142 downloader = download_images.ImageDownloader( 143 logger_to_use=MOCK_LOGGER, cmd_exec=mock_cmd_exec 144 ) 145 146 # Set os.path.exists to always return False and run uncompress. 147 mock_path_exists.return_value = False 148 self.assertRaises( 149 download_images.MissingImage, 150 downloader.UncompressImage, 151 test_chroot, 152 test_build_id, 153 ) 154 155 # Verify os.path.exists was called twice, with correct arguments. 156 self.assertEqual(mock_path_exists.call_count, 2) 157 mock_path_exists.assert_called_with( 158 RegexMatcher( 159 "/usr/local/home/chromeos/.*tmp/lumpy-release/" 160 "R36-5814.0.0/chromiumos_test_image.bin" 161 ) 162 ) 163 mock_path_exists.assert_any_call("/etc/cros_chroot_version") 164 165 # Verify RunCommand was called twice with correct arguments. 166 self.assertEqual(mock_cmd_exec.RunCommand.call_count, 2) 167 # Call 1, should have 2 arguments 168 self.assertEqual(len(mock_cmd_exec.RunCommand.call_args_list[0]), 2) 169 actual_arg = mock_cmd_exec.RunCommand.call_args_list[0][0] 170 expected_arg = ( 171 RegexMatcher( 172 "cd /usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0 ; " 173 "tar -Jxf chromiumos_test_image.tar.xz " 174 ), 175 ) 176 self.assertEqual(expected_arg, actual_arg) 177 # 2nd arg must be exception handler 178 except_handler_string = "RunCommandExceptionHandler.HandleException" 179 self.assertTrue( 180 except_handler_string 181 in repr(mock_cmd_exec.RunCommand.call_args_list[0][1]) 182 ) 183 184 # Call 2, should have 2 arguments 185 self.assertEqual(len(mock_cmd_exec.RunCommand.call_args_list[1]), 2) 186 actual_arg = mock_cmd_exec.RunCommand.call_args_list[1][0] 187 expected_arg = ( 188 RegexMatcher( 189 "cd /usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0 ; " 190 "rm -f chromiumos_test_image.bin " 191 ), 192 ) 193 self.assertEqual(expected_arg, actual_arg) 194 # 2nd arg must be empty 195 self.assertTrue( 196 "{}" in repr(mock_cmd_exec.RunCommand.call_args_list[1][1]) 197 ) 198 199 # Set os.path.exists to always return True (except for "inside chroot" 200 # check) and run uncompress. 201 mock_path_exists.reset_mock() 202 mock_cmd_exec.reset_mock() 203 mock_path_exists.side_effect = lambda x: x != "/etc/cros_chroot_version" 204 downloader.UncompressImage(test_chroot, test_build_id) 205 206 # Verify os.path.exists was called once, with correct arguments. 207 self.assertEqual(mock_path_exists.call_count, 2) 208 mock_path_exists.assert_called_with( 209 RegexMatcher( 210 "/usr/local/home/chromeos/.*tmp/lumpy-release/" 211 "R36-5814.0.0/chromiumos_test_image.bin" 212 ) 213 ) 214 mock_path_exists.assert_any_call("/etc/cros_chroot_version") 215 216 # Verify RunCommand was not called. 217 self.assertEqual(mock_cmd_exec.RunCommand.call_count, 0) 218 219 def test_run(self): 220 # Set test arguments 221 test_chroot = "/usr/local/home/chromeos" 222 test_build_id = "remote/lumpy/latest-dev" 223 test_empty_autotest_path = "" 224 test_empty_debug_path = "" 225 test_autotest_path = "/tmp/autotest" 226 test_debug_path = "/tmp/debug" 227 download_debug = True 228 229 # Set values to test/check. 230 self.called_download_image = False 231 self.called_uncompress_image = False 232 self.called_get_build_id = False 233 self.called_download_autotest_files = False 234 self.called_download_debug_file = False 235 236 # Define fake stub functions for Run to call 237 def FakeGetBuildID(unused_root, unused_xbuddy_label): 238 self.called_get_build_id = True 239 return "lumpy-release/R36-5814.0.0" 240 241 def GoodDownloadImage(root, build_id, image_path): 242 if root or build_id or image_path: 243 pass 244 self.called_download_image = True 245 return "chromiumos_test_image.bin" 246 247 def BadDownloadImage(root, build_id, image_path): 248 if root or build_id or image_path: 249 pass 250 self.called_download_image = True 251 raise download_images.MissingImage("Could not download image") 252 253 def FakeUncompressImage(root, build_id): 254 if root or build_id: 255 pass 256 self.called_uncompress_image = True 257 return 0 258 259 def FakeDownloadAutotestFiles(root, build_id): 260 if root or build_id: 261 pass 262 self.called_download_autotest_files = True 263 return "autotest" 264 265 def FakeDownloadDebugFile(root, build_id): 266 if root or build_id: 267 pass 268 self.called_download_debug_file = True 269 return "debug" 270 271 # Initialize downloader 272 downloader = download_images.ImageDownloader(logger_to_use=MOCK_LOGGER) 273 274 # Set downloader to call fake stubs. 275 downloader.GetBuildID = FakeGetBuildID 276 downloader.UncompressImage = FakeUncompressImage 277 downloader.DownloadImage = GoodDownloadImage 278 downloader.DownloadAutotestFiles = FakeDownloadAutotestFiles 279 downloader.DownloadDebugFile = FakeDownloadDebugFile 280 281 # Call Run. 282 image_path, autotest_path, debug_path = downloader.Run( 283 test_chroot, 284 test_build_id, 285 test_empty_autotest_path, 286 test_empty_debug_path, 287 download_debug, 288 ) 289 290 # Make sure it called both _DownloadImage and _UncompressImage 291 self.assertTrue(self.called_download_image) 292 self.assertTrue(self.called_uncompress_image) 293 # Make sure it called DownloadAutotestFiles 294 self.assertTrue(self.called_download_autotest_files) 295 # Make sure it called DownloadDebugFile 296 self.assertTrue(self.called_download_debug_file) 297 # Make sure it returned an image and autotest path returned from this call 298 self.assertTrue(image_path == "chromiumos_test_image.bin") 299 self.assertTrue(autotest_path == "autotest") 300 self.assertTrue(debug_path == "debug") 301 302 # Call Run with a non-empty autotest and debug path 303 self.called_download_autotest_files = False 304 self.called_download_debug_file = False 305 306 image_path, autotest_path, debug_path = downloader.Run( 307 test_chroot, 308 test_build_id, 309 test_autotest_path, 310 test_debug_path, 311 download_debug, 312 ) 313 314 # Verify that downloadAutotestFiles was not called 315 self.assertFalse(self.called_download_autotest_files) 316 # Make sure it returned the specified autotest path returned from this call 317 self.assertTrue(autotest_path == test_autotest_path) 318 # Make sure it returned the specified debug path returned from this call 319 self.assertTrue(debug_path == test_debug_path) 320 321 # Reset values; Now use fake stub that simulates DownloadImage failing. 322 self.called_download_image = False 323 self.called_uncompress_image = False 324 self.called_download_autotest_files = False 325 self.called_download_debug_file = False 326 downloader.DownloadImage = BadDownloadImage 327 328 # Call Run again. 329 self.assertRaises( 330 download_images.MissingImage, 331 downloader.Run, 332 test_chroot, 333 test_autotest_path, 334 test_debug_path, 335 test_build_id, 336 download_debug, 337 ) 338 339 # Verify that UncompressImage and downloadAutotestFiles were not called, 340 # since _DownloadImage "failed" 341 self.assertTrue(self.called_download_image) 342 self.assertFalse(self.called_uncompress_image) 343 self.assertFalse(self.called_download_autotest_files) 344 self.assertFalse(self.called_download_debug_file) 345 346 347if __name__ == "__main__": 348 unittest.main() 349