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