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 io
13import logging
14import unittest
15
16from devil import devil_env
17from devil.android import device_errors
18from devil.android import device_utils
19from devil.android import fastboot_utils
20from devil.android.sdk import fastboot
21from devil.utils import mock_calls
22
23with devil_env.SysPath(devil_env.PYMOCK_PATH):
24  import mock  # pylint: disable=import-error
25
26_BOARD = 'board_type'
27_SERIAL = '0123456789abcdef'
28_PARTITIONS = ['cache', 'userdata', 'system', 'bootloader', 'radio']
29_IMAGES = {
30    'cache': 'cache.img',
31    'userdata': 'userdata.img',
32    'system': 'system.img',
33    'bootloader': 'bootloader.img',
34    'radio': 'radio.img',
35}
36_VALID_FILES = [_BOARD + '.zip', 'android-info.txt']
37_INVALID_FILES = ['test.zip', 'android-info.txt']
38
39
40class MockFile(object):
41
42  def __init__(self, name='/tmp/some/file'):
43    self.file = mock.MagicMock(spec=file)
44    self.file.name = name
45
46  def __enter__(self):
47    return self.file
48
49  def __exit__(self, exc_type, exc_val, exc_tb):
50    pass
51
52  @property
53  def name(self):
54    return self.file.name
55
56
57def _FastbootWrapperMock(test_serial):
58  fastbooter = mock.Mock(spec=fastboot.Fastboot)
59  fastbooter.__str__ = mock.Mock(return_value=test_serial)
60  fastbooter.Devices.return_value = [test_serial]
61  return fastbooter
62
63
64def _DeviceUtilsMock(test_serial):
65  device = mock.Mock(spec=device_utils.DeviceUtils)
66  device.__str__ = mock.Mock(return_value=test_serial)
67  device.product_board = mock.Mock(return_value=_BOARD)
68  device.adb = mock.Mock()
69  return device
70
71
72class FastbootUtilsTest(mock_calls.TestCase):
73
74  def setUp(self):
75    self.device_utils_mock = _DeviceUtilsMock(_SERIAL)
76    self.fastboot_wrapper = _FastbootWrapperMock(_SERIAL)
77    self.fastboot = fastboot_utils.FastbootUtils(
78        self.device_utils_mock, fastbooter=self.fastboot_wrapper,
79        default_timeout=2, default_retries=0)
80    self.fastboot._board = _BOARD
81
82
83class FastbootUtilsInitTest(FastbootUtilsTest):
84
85  def testInitWithDeviceUtil(self):
86    f = fastboot_utils.FastbootUtils(self.device_utils_mock)
87    self.assertEqual(str(self.device_utils_mock), str(f._device))
88
89  def testInitWithMissing_fails(self):
90    with self.assertRaises(AttributeError):
91      fastboot_utils.FastbootUtils(None)
92    with self.assertRaises(AttributeError):
93      fastboot_utils.FastbootUtils('')
94
95
96class FastbootUtilsWaitForFastbootMode(FastbootUtilsTest):
97
98  # If this test fails by timing out after 1 second.
99  @mock.patch('time.sleep', mock.Mock())
100  def testWaitForFastbootMode(self):
101    self.fastboot.WaitForFastbootMode()
102
103
104class FastbootUtilsEnableFastbootMode(FastbootUtilsTest):
105
106  def testEnableFastbootMode(self):
107    with self.assertCalls(
108        self.call.fastboot._device.EnableRoot(),
109        self.call.fastboot._device.adb.Reboot(to_bootloader=True),
110        self.call.fastboot.WaitForFastbootMode()):
111      self.fastboot.EnableFastbootMode()
112
113
114class FastbootUtilsReboot(FastbootUtilsTest):
115
116  def testReboot_bootloader(self):
117    with self.assertCalls(
118        self.call.fastboot.fastboot.RebootBootloader(),
119        self.call.fastboot.WaitForFastbootMode()):
120      self.fastboot.Reboot(bootloader=True)
121
122  def testReboot_normal(self):
123    with self.assertCalls(
124        self.call.fastboot.fastboot.Reboot(),
125        self.call.fastboot._device.WaitUntilFullyBooted(timeout=mock.ANY)):
126      self.fastboot.Reboot()
127
128
129class FastbootUtilsFlashPartitions(FastbootUtilsTest):
130
131  def testFlashPartitions_wipe(self):
132    with self.assertCalls(
133        (self.call.fastboot._VerifyBoard('test'), True),
134        (self.call.fastboot._FindAndVerifyPartitionsAndImages(
135            _PARTITIONS, 'test'), _IMAGES),
136        (self.call.fastboot.fastboot.Flash('cache', 'cache.img')),
137        (self.call.fastboot.fastboot.Flash('userdata', 'userdata.img')),
138        (self.call.fastboot.fastboot.Flash('system', 'system.img')),
139        (self.call.fastboot.fastboot.Flash('bootloader', 'bootloader.img')),
140        (self.call.fastboot.Reboot(bootloader=True)),
141        (self.call.fastboot.fastboot.Flash('radio', 'radio.img')),
142        (self.call.fastboot.Reboot(bootloader=True))):
143      self.fastboot._FlashPartitions(_PARTITIONS, 'test', wipe=True)
144
145  def testFlashPartitions_noWipe(self):
146    with self.assertCalls(
147        (self.call.fastboot._VerifyBoard('test'), True),
148        (self.call.fastboot._FindAndVerifyPartitionsAndImages(
149            _PARTITIONS, 'test'), _IMAGES),
150        (self.call.fastboot.fastboot.Flash('system', 'system.img')),
151        (self.call.fastboot.fastboot.Flash('bootloader', 'bootloader.img')),
152        (self.call.fastboot.Reboot(bootloader=True)),
153        (self.call.fastboot.fastboot.Flash('radio', 'radio.img')),
154        (self.call.fastboot.Reboot(bootloader=True))):
155      self.fastboot._FlashPartitions(_PARTITIONS, 'test')
156
157
158class FastbootUtilsFastbootMode(FastbootUtilsTest):
159
160  def testFastbootMode_good(self):
161    with self.assertCalls(
162        self.call.fastboot.EnableFastbootMode(),
163        self.call.fastboot.fastboot.SetOemOffModeCharge(False),
164        self.call.fastboot.fastboot.SetOemOffModeCharge(True),
165        self.call.fastboot.Reboot()):
166      with self.fastboot.FastbootMode() as fbm:
167        self.assertEqual(self.fastboot, fbm)
168
169  def testFastbootMode_exception(self):
170    with self.assertCalls(
171        self.call.fastboot.EnableFastbootMode(),
172        self.call.fastboot.fastboot.SetOemOffModeCharge(False),
173        self.call.fastboot.fastboot.SetOemOffModeCharge(True),
174        self.call.fastboot.Reboot()):
175      with self.assertRaises(NotImplementedError):
176        with self.fastboot.FastbootMode() as fbm:
177          self.assertEqual(self.fastboot, fbm)
178          raise NotImplementedError
179
180  def testFastbootMode_exceptionInEnableFastboot(self):
181    self.fastboot.EnableFastbootMode = mock.Mock()
182    self.fastboot.EnableFastbootMode.side_effect = NotImplementedError
183    with self.assertRaises(NotImplementedError):
184      with self.fastboot.FastbootMode():
185        pass
186
187
188class FastbootUtilsVerifyBoard(FastbootUtilsTest):
189
190  def testVerifyBoard_bothValid(self):
191    mock_file = io.StringIO(u'require board=%s\n' % _BOARD)
192    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
193      with mock.patch('os.listdir', return_value=_VALID_FILES):
194        self.assertTrue(self.fastboot._VerifyBoard('test'))
195
196  def testVerifyBoard_BothNotValid(self):
197    mock_file = io.StringIO(u'abc')
198    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
199      with mock.patch('os.listdir', return_value=_INVALID_FILES):
200        self.assertFalse(self.assertFalse(self.fastboot._VerifyBoard('test')))
201
202  def testVerifyBoard_FileNotFoundZipValid(self):
203    with mock.patch('os.listdir', return_value=[_BOARD + '.zip']):
204      self.assertTrue(self.fastboot._VerifyBoard('test'))
205
206  def testVerifyBoard_ZipNotFoundFileValid(self):
207    mock_file = io.StringIO(u'require board=%s\n' % _BOARD)
208    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
209      with mock.patch('os.listdir', return_value=['android-info.txt']):
210        self.assertTrue(self.fastboot._VerifyBoard('test'))
211
212  def testVerifyBoard_zipNotValidFileIs(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=_INVALID_FILES):
216        self.assertTrue(self.fastboot._VerifyBoard('test'))
217
218  def testVerifyBoard_fileNotValidZipIs(self):
219    mock_file = io.StringIO(u'require board=WrongBoard')
220    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
221      with mock.patch('os.listdir', return_value=_VALID_FILES):
222        self.assertFalse(self.fastboot._VerifyBoard('test'))
223
224  def testVerifyBoard_noBoardInFileValidZip(self):
225    mock_file = io.StringIO(u'Regex wont match')
226    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
227      with mock.patch('os.listdir', return_value=_VALID_FILES):
228        self.assertTrue(self.fastboot._VerifyBoard('test'))
229
230  def testVerifyBoard_noBoardInFileInvalidZip(self):
231    mock_file = io.StringIO(u'Regex wont match')
232    with mock.patch('__builtin__.open', return_value=mock_file, create=True):
233      with mock.patch('os.listdir', return_value=_INVALID_FILES):
234        self.assertFalse(self.fastboot._VerifyBoard('test'))
235
236
237class FastbootUtilsFindAndVerifyPartitionsAndImages(FastbootUtilsTest):
238
239  def testFindAndVerifyPartitionsAndImages_valid(self):
240    PARTITIONS = [
241        'bootloader', 'radio', 'boot', 'recovery', 'system', 'userdata', 'cache'
242    ]
243    files = [
244        'bootloader-test-.img',
245        'radio123.img',
246        'boot.img',
247        'recovery.img',
248        'system.img',
249        'userdata.img',
250        'cache.img'
251    ]
252    return_check = {
253      'bootloader': 'test/bootloader-test-.img',
254      'radio': 'test/radio123.img',
255      'boot': 'test/boot.img',
256      'recovery': 'test/recovery.img',
257      'system': 'test/system.img',
258      'userdata': 'test/userdata.img',
259      'cache': 'test/cache.img',
260    }
261
262    with mock.patch('os.listdir', return_value=files):
263      return_value = self.fastboot._FindAndVerifyPartitionsAndImages(
264          PARTITIONS, 'test')
265      self.assertDictEqual(return_value, return_check)
266
267  def testFindAndVerifyPartitionsAndImages_badPartition(self):
268    with mock.patch('os.listdir', return_value=['test']):
269      with self.assertRaises(KeyError):
270        self.fastboot._FindAndVerifyPartitionsAndImages(['test'], 'test')
271
272  def testFindAndVerifyPartitionsAndImages_noFile(self):
273    with mock.patch('os.listdir', return_value=['test']):
274      with self.assertRaises(device_errors.FastbootCommandFailedError):
275        self.fastboot._FindAndVerifyPartitionsAndImages(['cache'], 'test')
276
277
278if __name__ == '__main__':
279  logging.getLogger().setLevel(logging.DEBUG)
280  unittest.main(verbosity=2)
281