download_images_unittest.py revision 3fb928efb42405e386da41adf4bd674e86e23022 
1#!/usr/bin/env python2
2#
3# Copyright 2014 Google Inc.  All Rights Reserved
4"""Download image unittest."""
5
6from __future__ import print_function
7
8import os
9import mock
10import unittest
11
12import download_images
13from cros_utils import command_executer
14from cros_utils import logger
15
16import test_flag
17
18MOCK_LOGGER = logger.GetLogger(log_dir='', mock=True)
19
20
21class 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
250if __name__ == '__main__':
251  unittest.main()
252