1#!/usr/bin/env python
2# Copyright 2015 The Chromium Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""
7Unit tests for the contents of fastboot_utils.py
8"""
9
10# pylint: disable=protected-access,unused-argument
11
12import collections
13import io
14import logging
15import unittest
16
17from devil import devil_env
18from devil.android import device_errors
19from devil.android import device_utils
20from devil.android import fastboot_utils
21from devil.android.sdk import fastboot
22from devil.utils import mock_calls
23
24with devil_env.SysPath(devil_env.PYMOCK_PATH):
25  import mock  # pylint: disable=import-error
26
27_BOARD = 'board_type'
28_SERIAL = '0123456789abcdef'
29_PARTITIONS = [
30    'bootloader', 'radio', 'boot', 'recovery', 'system', 'userdata', 'cache']
31_IMAGES = collections.OrderedDict([
32    ('bootloader', 'bootloader.img'),
33    ('radio', 'radio.img'),
34    ('boot', 'boot.img'),
35    ('recovery', 'recovery.img'),
36    ('system', 'system.img'),
37    ('userdata', 'userdata.img'),
38    ('cache', 'cache.img')
39])
40_VALID_FILES = [_BOARD + '.zip', 'android-info.txt']
41_INVALID_FILES = ['test.zip', 'android-info.txt']
42
43
44class MockFile(object):
45
46  def __init__(self, name='/tmp/some/file'):
47    self.file = mock.MagicMock(spec=file)
48    self.file.name = name
49
50  def __enter__(self):
51    return self.file
52
53  def __exit__(self, exc_type, exc_val, exc_tb):
54    pass
55
56  @property
57  def name(self):
58    return self.file.name
59
60
61def _FastbootWrapperMock(test_serial):
62  fastbooter = mock.Mock(spec=fastboot.Fastboot)
63  fastbooter.__str__ = mock.Mock(return_value=test_serial)
64  fastbooter.Devices.return_value = [test_serial]
65  return fastbooter
66
67
68def _DeviceUtilsMock(test_serial):
69  device = mock.Mock(spec=device_utils.DeviceUtils)
70  device.__str__ = mock.Mock(return_value=test_serial)
71  device.product_board = mock.Mock(return_value=_BOARD)
72  device.adb = mock.Mock()
73  return device
74
75
76class FastbootUtilsTest(mock_calls.TestCase):
77
78  def setUp(self):
79    self.device_utils_mock = _DeviceUtilsMock(_SERIAL)
80    self.fastboot_wrapper = _FastbootWrapperMock(_SERIAL)
81    self.fastboot = fastboot_utils.FastbootUtils(
82        self.device_utils_mock, fastbooter=self.fastboot_wrapper,
83        default_timeout=2, default_retries=0)
84    self.fastboot._board = _BOARD
85
86
87class FastbootUtilsInitTest(FastbootUtilsTest):
88
89  def testInitWithDeviceUtil(self):
90    f = fastboot_utils.FastbootUtils(self.device_utils_mock)
91    self.assertEqual(str(self.device_utils_mock), str(f._device))
92
93  def testInitWithMissing_fails(self):
94    with self.assertRaises(AttributeError):
95      fastboot_utils.FastbootUtils(None)
96    with self.assertRaises(AttributeError):
97      fastboot_utils.FastbootUtils('')
98
99  def testPartitionOrdering(self):
100    parts = ['bootloader', 'radio', 'boot', 'recovery', 'system', 'userdata',
101             'cache', 'vendor']
102    self.assertListEqual(fastboot_utils.ALL_PARTITIONS, parts)
103
104
105class FastbootUtilsWaitForFastbootMode(FastbootUtilsTest):
106
107  # If this test fails by timing out after 1 second.
108  @mock.patch('time.sleep', mock.Mock())
109  def testWaitForFastbootMode(self):
110    self.fastboot.WaitForFastbootMode()
111
112
113class FastbootUtilsEnableFastbootMode(FastbootUtilsTest):
114
115  def testEnableFastbootMode(self):
116    with self.assertCalls(
117        self.call.fastboot._device.EnableRoot(),
118        self.call.fastboot._device.adb.Reboot(to_bootloader=True),
119        self.call.fastboot.WaitForFastbootMode()):
120      self.fastboot.EnableFastbootMode()
121
122
123class FastbootUtilsReboot(FastbootUtilsTest):
124
125  def testReboot_bootloader(self):
126    with self.assertCalls(
127        self.call.fastboot.fastboot.RebootBootloader(),
128        self.call.fastboot.WaitForFastbootMode()):
129      self.fastboot.Reboot(bootloader=True)
130
131  def testReboot_normal(self):
132    with self.assertCalls(
133        self.call.fastboot.fastboot.Reboot(),
134        self.call.fastboot._device.WaitUntilFullyBooted(timeout=mock.ANY)):
135      self.fastboot.Reboot()
136
137
138class FastbootUtilsFlashPartitions(FastbootUtilsTest):
139
140  def testFlashPartitions_wipe(self):
141    with self.assertCalls(
142        (self.call.fastboot._VerifyBoard('test'), True),
143        (mock.call.devil.android.fastboot_utils.
144            _FindAndVerifyPartitionsAndImages(_PARTITIONS, 'test'), _IMAGES),
145        (self.call.fastboot.fastboot.Flash('bootloader', 'bootloader.img')),
146        (self.call.fastboot.Reboot(bootloader=True)),
147        (self.call.fastboot.fastboot.Flash('radio', 'radio.img')),
148        (self.call.fastboot.Reboot(bootloader=True)),
149        (self.call.fastboot.fastboot.Flash('boot', 'boot.img')),
150        (self.call.fastboot.fastboot.Flash('recovery', 'recovery.img')),
151        (self.call.fastboot.fastboot.Flash('system', 'system.img')),
152        (self.call.fastboot.fastboot.Flash('userdata', 'userdata.img')),
153        (self.call.fastboot.fastboot.Flash('cache', 'cache.img'))):
154      self.fastboot._FlashPartitions(_PARTITIONS, 'test', wipe=True)
155
156  def testFlashPartitions_noWipe(self):
157    with self.assertCalls(
158        (self.call.fastboot._VerifyBoard('test'), True),
159        (mock.call.devil.android.fastboot_utils.
160            _FindAndVerifyPartitionsAndImages(_PARTITIONS, 'test'), _IMAGES),
161        (self.call.fastboot.fastboot.Flash('bootloader', 'bootloader.img')),
162        (self.call.fastboot.Reboot(bootloader=True)),
163        (self.call.fastboot.fastboot.Flash('radio', 'radio.img')),
164        (self.call.fastboot.Reboot(bootloader=True)),
165        (self.call.fastboot.fastboot.Flash('boot', 'boot.img')),
166        (self.call.fastboot.fastboot.Flash('recovery', 'recovery.img')),
167        (self.call.fastboot.fastboot.Flash('system', 'system.img'))):
168      self.fastboot._FlashPartitions(_PARTITIONS, 'test')
169
170
171class FastbootUtilsFastbootMode(FastbootUtilsTest):
172
173  def testFastbootMode_goodWait(self):
174    with self.assertCalls(
175        self.call.fastboot.EnableFastbootMode(),
176        self.call.fastboot.fastboot.SetOemOffModeCharge(False),
177        self.call.fastboot.fastboot.SetOemOffModeCharge(True),
178        self.call.fastboot.Reboot(wait_for_reboot=True)):
179      with self.fastboot.FastbootMode() as fbm:
180        self.assertEqual(self.fastboot, fbm)
181
182  def testFastbootMode_goodNoWait(self):
183    with self.assertCalls(
184        self.call.fastboot.EnableFastbootMode(),
185        self.call.fastboot.fastboot.SetOemOffModeCharge(False),
186        self.call.fastboot.fastboot.SetOemOffModeCharge(True),
187        self.call.fastboot.Reboot(wait_for_reboot=False)):
188      with self.fastboot.FastbootMode(wait_for_reboot=False) as fbm:
189        self.assertEqual(self.fastboot, fbm)
190
191  def testFastbootMode_exception(self):
192    with self.assertCalls(
193        self.call.fastboot.EnableFastbootMode(),
194        self.call.fastboot.fastboot.SetOemOffModeCharge(False),
195        self.call.fastboot.fastboot.SetOemOffModeCharge(True),
196        self.call.fastboot.Reboot(wait_for_reboot=True)):
197      with self.assertRaises(NotImplementedError):
198        with self.fastboot.FastbootMode() as fbm:
199          self.assertEqual(self.fastboot, fbm)
200          raise NotImplementedError
201
202  def testFastbootMode_exceptionInEnableFastboot(self):
203    self.fastboot.EnableFastbootMode = mock.Mock()
204    self.fastboot.EnableFastbootMode.side_effect = NotImplementedError
205    with self.assertRaises(NotImplementedError):
206      with self.fastboot.FastbootMode():
207        pass
208
209
210class FastbootUtilsVerifyBoard(FastbootUtilsTest):
211
212  def testVerifyBoard_bothValid(self):
213    mock_file = io.StringIO(u'require board=%s\n' % _BOARD)
214    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
215      with mock.patch('os.listdir', return_value=_VALID_FILES):
216        self.assertTrue(self.fastboot._VerifyBoard('test'))
217
218  def testVerifyBoard_BothNotValid(self):
219    mock_file = io.StringIO(u'abc')
220    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
221      with mock.patch('os.listdir', return_value=_INVALID_FILES):
222        self.assertFalse(self.assertFalse(self.fastboot._VerifyBoard('test')))
223
224  def testVerifyBoard_FileNotFoundZipValid(self):
225    with mock.patch('os.listdir', return_value=[_BOARD + '.zip']):
226      self.assertTrue(self.fastboot._VerifyBoard('test'))
227
228  def testVerifyBoard_ZipNotFoundFileValid(self):
229    mock_file = io.StringIO(u'require board=%s\n' % _BOARD)
230    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
231      with mock.patch('os.listdir', return_value=['android-info.txt']):
232        self.assertTrue(self.fastboot._VerifyBoard('test'))
233
234  def testVerifyBoard_zipNotValidFileIs(self):
235    mock_file = io.StringIO(u'require board=%s\n' % _BOARD)
236    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
237      with mock.patch('os.listdir', return_value=_INVALID_FILES):
238        self.assertTrue(self.fastboot._VerifyBoard('test'))
239
240  def testVerifyBoard_fileNotValidZipIs(self):
241    mock_file = io.StringIO(u'require board=WrongBoard')
242    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
243      with mock.patch('os.listdir', return_value=_VALID_FILES):
244        self.assertFalse(self.fastboot._VerifyBoard('test'))
245
246  def testVerifyBoard_noBoardInFileValidZip(self):
247    mock_file = io.StringIO(u'Regex wont match')
248    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
249      with mock.patch('os.listdir', return_value=_VALID_FILES):
250        self.assertTrue(self.fastboot._VerifyBoard('test'))
251
252  def testVerifyBoard_noBoardInFileInvalidZip(self):
253    mock_file = io.StringIO(u'Regex wont match')
254    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
255      with mock.patch('os.listdir', return_value=_INVALID_FILES):
256        self.assertFalse(self.fastboot._VerifyBoard('test'))
257
258
259class FastbootUtilsFindAndVerifyPartitionsAndImages(FastbootUtilsTest):
260
261  def testFindAndVerifyPartitionsAndImages_validNoVendor(self):
262    PARTITIONS = [
263        'bootloader', 'radio', 'boot', 'recovery', 'system', 'userdata',
264        'cache', 'vendor'
265    ]
266    files = [
267        'bootloader-test-.img',
268        'radio123.img',
269        'boot.img',
270        'recovery.img',
271        'system.img',
272        'userdata.img',
273        'cache.img'
274    ]
275    img_check = collections.OrderedDict([
276      ('bootloader', 'test/bootloader-test-.img'),
277      ('radio', 'test/radio123.img'),
278      ('boot', 'test/boot.img'),
279      ('recovery', 'test/recovery.img'),
280      ('system', 'test/system.img'),
281      ('userdata', 'test/userdata.img'),
282      ('cache', 'test/cache.img'),
283    ])
284    parts_check = [
285        'bootloader', 'radio', 'boot', 'recovery', 'system', 'userdata',
286        'cache'
287    ]
288    with mock.patch('os.listdir', return_value=files):
289      imgs = fastboot_utils._FindAndVerifyPartitionsAndImages(
290          PARTITIONS, 'test')
291      parts = imgs.keys()
292      self.assertDictEqual(imgs, img_check)
293      self.assertListEqual(parts, parts_check)
294
295  def testFindAndVerifyPartitionsAndImages_validVendor(self):
296    PARTITIONS = [
297        'bootloader', 'radio', 'boot', 'recovery', 'system', 'userdata',
298        'cache', 'vendor'
299    ]
300    files = [
301        'bootloader-test-.img',
302        'radio123.img',
303        'boot.img',
304        'recovery.img',
305        'system.img',
306        'userdata.img',
307        'cache.img',
308        'vendor.img'
309    ]
310    img_check = {
311      'bootloader': 'test/bootloader-test-.img',
312      'radio': 'test/radio123.img',
313      'boot': 'test/boot.img',
314      'recovery': 'test/recovery.img',
315      'system': 'test/system.img',
316      'userdata': 'test/userdata.img',
317      'cache': 'test/cache.img',
318      'vendor': 'test/vendor.img',
319    }
320    parts_check = [
321        'bootloader', 'radio', 'boot', 'recovery', 'system', 'userdata',
322        'cache', 'vendor'
323    ]
324
325    with mock.patch('os.listdir', return_value=files):
326      imgs = fastboot_utils._FindAndVerifyPartitionsAndImages(
327          PARTITIONS, 'test')
328      parts = imgs.keys()
329      self.assertDictEqual(imgs, img_check)
330      self.assertListEqual(parts, parts_check)
331
332  def testFindAndVerifyPartitionsAndImages_badPartition(self):
333    with mock.patch('os.listdir', return_value=['test']):
334      with self.assertRaises(KeyError):
335        fastboot_utils._FindAndVerifyPartitionsAndImages(['test'], 'test')
336
337  def testFindAndVerifyPartitionsAndImages_noFile(self):
338    with mock.patch('os.listdir', return_value=['test']):
339      with self.assertRaises(device_errors.FastbootCommandFailedError):
340        fastboot_utils._FindAndVerifyPartitionsAndImages(['cache'], 'test')
341
342
343class FastbootUtilsFlashDevice(FastbootUtilsTest):
344
345  def testFlashDevice_wipe(self):
346    with self.assertCalls(
347        self.call.fastboot.EnableFastbootMode(),
348        self.call.fastboot.fastboot.SetOemOffModeCharge(False),
349        self.call.fastboot._FlashPartitions(mock.ANY, 'test', wipe=True),
350        self.call.fastboot.fastboot.SetOemOffModeCharge(True),
351        self.call.fastboot.Reboot(wait_for_reboot=False)):
352      self.fastboot.FlashDevice('test', wipe=True)
353
354  def testFlashDevice_noWipe(self):
355    with self.assertCalls(
356        self.call.fastboot.EnableFastbootMode(),
357        self.call.fastboot.fastboot.SetOemOffModeCharge(False),
358        self.call.fastboot._FlashPartitions(mock.ANY, 'test', wipe=False),
359        self.call.fastboot.fastboot.SetOemOffModeCharge(True),
360        self.call.fastboot.Reboot(wait_for_reboot=True)):
361      self.fastboot.FlashDevice('test', wipe=False)
362
363  def testFlashDevice_partitions(self):
364    with self.assertCalls(
365        self.call.fastboot.EnableFastbootMode(),
366        self.call.fastboot.fastboot.SetOemOffModeCharge(False),
367        self.call.fastboot._FlashPartitions(['boot'], 'test', wipe=False),
368        self.call.fastboot.fastboot.SetOemOffModeCharge(True),
369        self.call.fastboot.Reboot(wait_for_reboot=True)):
370      self.fastboot.FlashDevice('test', partitions=['boot'], wipe=False)
371
372
373if __name__ == '__main__':
374  logging.getLogger().setLevel(logging.DEBUG)
375  unittest.main(verbosity=2)
376