Home | History | Annotate | Download | only in crosperf
      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
    130         repr(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