download_images_unittest.py revision c7ed2bca6f5dbf0cc6ee7cd1e7d30ea3119adbe8 
1#!/usr/bin/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
30  @mock.patch.object(os, 'makedirs')
31  @mock.patch.object(os.path, 'exists')
32  def test_download_image(self, mock_path_exists, mock_mkdirs):
33
34    # Set mock and test values.
35    mock_cmd_exec = mock.Mock(spec=command_executer.CommandExecuter)
36    test_chroot = '/usr/local/home/chromeos'
37    test_build_id = 'lumpy-release/R36-5814.0.0'
38    image_path = ('gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz'
39                  % test_build_id)
40
41    downloader = download_images.ImageDownloader(logger_to_use=MOCK_LOGGER,
42                                                 cmd_exec=mock_cmd_exec)
43
44    # Set os.path.exists to always return False and run downloader
45    mock_path_exists.return_value = False
46    test_flag.SetTestMode(True)
47    self.assertRaises(download_images.MissingImage, downloader.DownloadImage,
48                      test_chroot, test_build_id, image_path)
49
50    # Verify os.path.exists was called twice, with proper arguments.
51    self.assertEqual(mock_path_exists.call_count, 2)
52    mock_path_exists.assert_called_with(
53        '/usr/local/home/chromeos/chroot/tmp/lumpy-release/'
54        'R36-5814.0.0/chromiumos_test_image.bin')
55    mock_path_exists.assert_any_call(
56        '/usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0')
57
58    # Verify we called os.mkdirs
59    self.assertEqual(mock_mkdirs.call_count, 1)
60    mock_mkdirs.assert_called_with(
61        '/usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0')
62
63    # Verify we called ChrootRunCommand once, with proper arguments.
64    self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 1)
65    mock_cmd_exec.ChrootRunCommand.assert_called_with(
66        '/usr/local/home/chromeos', 'gsutil cp '
67        'gs://chromeos-image-archive/lumpy-release/R36-5814.0.0/'
68        'chromiumos_test_image.tar.xz'
69        ' /tmp/lumpy-release/R36-5814.0.0')
70
71    # Reset the velues in the mocks; set os.path.exists to always return True.
72    mock_path_exists.reset_mock()
73    mock_cmd_exec.reset_mock()
74    mock_path_exists.return_value = True
75
76    # Run downloader
77    downloader.DownloadImage(test_chroot, test_build_id, image_path)
78
79    # Verify os.path.exists was called twice, with proper arguments.
80    self.assertEqual(mock_path_exists.call_count, 2)
81    mock_path_exists.assert_called_with(
82        '/usr/local/home/chromeos/chroot/tmp/lumpy-release/'
83        'R36-5814.0.0/chromiumos_test_image.bin')
84    mock_path_exists.assert_any_call(
85        '/usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0')
86
87    # Verify we made no RunCommand or ChrootRunCommand calls (since
88    # os.path.exists returned True, there was no work do be done).
89    self.assertEqual(mock_cmd_exec.RunCommand.call_count, 0)
90    self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 0)
91
92  @mock.patch.object(os.path, 'exists')
93  def test_uncompress_image(self, mock_path_exists):
94
95    # set mock and test values.
96    mock_cmd_exec = mock.Mock(spec=command_executer.CommandExecuter)
97    test_chroot = '/usr/local/home/chromeos'
98    test_build_id = 'lumpy-release/R36-5814.0.0'
99
100    downloader = download_images.ImageDownloader(logger_to_use=MOCK_LOGGER,
101                                                 cmd_exec=mock_cmd_exec)
102
103    # Set os.path.exists to always return False and run uncompress.
104    mock_path_exists.return_value = False
105    self.assertRaises(download_images.MissingImage, downloader.UncompressImage,
106                      test_chroot, test_build_id)
107
108    # Verify os.path.exists was called once, with correct arguments.
109    self.assertEqual(mock_path_exists.call_count, 1)
110    mock_path_exists.assert_called_with(
111        '/usr/local/home/chromeos/chroot/tmp/lumpy-release/'
112        'R36-5814.0.0/chromiumos_test_image.bin')
113
114    # Verify ChrootRunCommand was called, with correct arguments.
115    self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 1)
116    mock_cmd_exec.ChrootRunCommand.assert_called_with(
117        '/usr/local/home/chromeos', 'cd /tmp/lumpy-release/R36-5814.0.0 ;unxz '
118        'chromiumos_test_image.tar.xz; tar -xvf chromiumos_test_image.tar')
119
120    # Set os.path.exists to always return False and run uncompress.
121    mock_path_exists.reset_mock()
122    mock_cmd_exec.reset_mock()
123    mock_path_exists.return_value = True
124    downloader.UncompressImage(test_chroot, test_build_id)
125
126    # Verify os.path.exists was called once, with correct arguments.
127    self.assertEqual(mock_path_exists.call_count, 1)
128    mock_path_exists.assert_called_with(
129        '/usr/local/home/chromeos/chroot/tmp/lumpy-release/'
130        'R36-5814.0.0/chromiumos_test_image.bin')
131
132    # Verify ChrootRunCommand was not called.
133    self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 0)
134
135  def test_run(self):
136
137    # Set test arguments
138    test_chroot = '/usr/local/home/chromeos'
139    test_build_id = 'remote/lumpy/latest-dev'
140
141    # Set values to test/check.
142    self.called_download_image = False
143    self.called_uncompress_image = False
144    self.called_get_build_id = False
145
146    # Define fake stub functions for Run to call
147    def FakeGetBuildID(unused_root, unused_xbuddy_label):
148      self.called_get_build_id = True
149      return 'lumpy-release/R36-5814.0.0'
150
151    def GoodDownloadImage(root, build_id, image_path):
152      if root or build_id or image_path:
153        pass
154      self.called_download_image = True
155      return 'chromiumos_test_image.bin'
156
157    def BadDownloadImage(root, build_id, image_path):
158      if root or build_id or image_path:
159        pass
160      self.called_download_image = True
161      raise download_images.MissingImage('Could not download image')
162
163    def FakeUncompressImage(root, build_id):
164      if root or build_id:
165        pass
166      self.called_uncompress_image = True
167      return 0
168
169    # Initialize downloader
170    downloader = download_images.ImageDownloader(logger_to_use=MOCK_LOGGER)
171
172    # Set downloader to call fake stubs.
173    downloader.GetBuildID = FakeGetBuildID
174    downloader.UncompressImage = FakeUncompressImage
175    downloader.DownloadImage = GoodDownloadImage
176
177    # Call Run.
178    downloader.Run(test_chroot, test_build_id)
179
180    # Make sure it called both _DownloadImage and _UncompressImage
181    self.assertTrue(self.called_download_image)
182    self.assertTrue(self.called_uncompress_image)
183
184    # Reset values; Now use fake stub that simulates DownloadImage failing.
185    self.called_download_image = False
186    self.called_uncompress_image = False
187    downloader.DownloadImage = BadDownloadImage
188
189    # Call Run again.
190    self.assertRaises(download_images.MissingImage, downloader.Run, test_chroot,
191                      test_build_id)
192
193    # Verify that UncompressImage was not called, since _DownloadImage "failed"
194    self.assertTrue(self.called_download_image)
195    self.assertFalse(self.called_uncompress_image)
196
197
198if __name__ == '__main__':
199  unittest.main()
200