blob: bc511b7b20819031fa78ce29cfa8a04997256cf7 [file] [log] [blame]
Mike Frysingere58c0e22017-10-04 15:43:30 -04001# -*- coding: utf-8 -*-
David Pursellf1d16a62015-03-25 13:31:04 -07002# Copyright 2015 The Chromium OS 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"""Unit tests for the flash module."""
7
8from __future__ import print_function
9
David Pursellf1d16a62015-03-25 13:31:04 -070010import os
11
Mike Frysinger6db648e2018-07-24 19:57:58 -040012import mock
13
David Pursellf1d16a62015-03-25 13:31:04 -070014from chromite.cli import flash
Sanika Kulkarnia8c4e3a2019-09-20 16:47:25 -070015from chromite.lib import auto_updater_transfer
David Pursellf1d16a62015-03-25 13:31:04 -070016from chromite.lib import commandline
17from chromite.lib import cros_build_lib
Ralph Nathan9b997232015-05-15 13:13:12 -070018from chromite.lib import cros_logging as logging
David Pursellf1d16a62015-03-25 13:31:04 -070019from chromite.lib import cros_test_lib
20from chromite.lib import dev_server_wrapper
Bertrand SIMONNET56f773d2015-05-04 14:02:39 -070021from chromite.lib import osutils
David Pursellf1d16a62015-03-25 13:31:04 -070022from chromite.lib import partial_mock
23from chromite.lib import remote_access
24
Amin Hassanic0f06fa2019-01-28 15:24:47 -080025from chromite.lib.paygen import paygen_payload_lib
26from chromite.lib.paygen import paygen_stateful_payload_lib
27
David Pursellf1d16a62015-03-25 13:31:04 -070028
29class RemoteDeviceUpdaterMock(partial_mock.PartialCmdMock):
30 """Mock out RemoteDeviceUpdater."""
Amin Hassani9800d432019-07-24 14:23:39 -070031 TARGET = 'chromite.lib.auto_updater.ChromiumOSUpdater'
xixuane851dfb2016-05-02 18:02:37 -070032 ATTRS = ('UpdateStateful', 'UpdateRootfs', 'SetupRootfsUpdate',
Amin Hassanif85be122019-09-19 17:10:34 -070033 'RebootAndVerify', 'ResolveAPPIDMismatchIfAny')
David Pursellf1d16a62015-03-25 13:31:04 -070034
35 def __init__(self):
36 partial_mock.PartialCmdMock.__init__(self)
37
38 def UpdateStateful(self, _inst, *_args, **_kwargs):
39 """Mock out UpdateStateful."""
40
41 def UpdateRootfs(self, _inst, *_args, **_kwargs):
42 """Mock out UpdateRootfs."""
43
44 def SetupRootfsUpdate(self, _inst, *_args, **_kwargs):
45 """Mock out SetupRootfsUpdate."""
46
xixuane851dfb2016-05-02 18:02:37 -070047 def RebootAndVerify(self, _inst, *_args, **_kwargs):
48 """Mock out RebootAndVerify."""
David Pursellf1d16a62015-03-25 13:31:04 -070049
Amin Hassanif85be122019-09-19 17:10:34 -070050 def ResolveAPPIDMismatchIfAny(self, _inst, *_args, **_kwargs):
51 """Mock out ResolveAPPIDMismatchIfAny."""
David Pursellf1d16a62015-03-25 13:31:04 -070052
Achuith Bhandarkar4d4275f2019-10-01 17:07:23 +020053
David Pursellf1d16a62015-03-25 13:31:04 -070054class RemoteDeviceUpdaterTest(cros_test_lib.MockTempDirTestCase):
55 """Test the flow of flash.Flash() with RemoteDeviceUpdater."""
56
57 IMAGE = '/path/to/image'
58 DEVICE = commandline.Device(scheme=commandline.DEVICE_SCHEME_SSH,
59 hostname='1.1.1.1')
60
61 def setUp(self):
62 """Patches objects."""
63 self.updater_mock = self.StartPatcher(RemoteDeviceUpdaterMock())
64 self.PatchObject(dev_server_wrapper, 'GenerateXbuddyRequest',
65 return_value='xbuddy/local/latest')
David Pursellf1d16a62015-03-25 13:31:04 -070066 self.PatchObject(dev_server_wrapper, 'GetImagePathWithXbuddy',
Gilad Arnolde62ec902015-04-24 14:41:02 -070067 return_value=('taco-paladin/R36/chromiumos_test_image.bin',
68 'remote/taco-paladin/R36/test'))
Amin Hassanic0f06fa2019-01-28 15:24:47 -080069 self.PatchObject(paygen_payload_lib, 'GenerateUpdatePayload')
70 self.PatchObject(paygen_stateful_payload_lib, 'GenerateStatefulPayload')
David Pursellf1d16a62015-03-25 13:31:04 -070071 self.PatchObject(remote_access, 'CHECK_INTERVAL', new=0)
72 self.PatchObject(remote_access, 'ChromiumOSDevice')
73
74 def testUpdateAll(self):
75 """Tests that update methods are called correctly."""
David Pursellf1d16a62015-03-25 13:31:04 -070076 with mock.patch('os.path.exists', return_value=True):
Bertrand SIMONNETb34a98b2015-04-22 14:30:04 -070077 flash.Flash(self.DEVICE, self.IMAGE)
78 self.assertTrue(self.updater_mock.patched['UpdateStateful'].called)
79 self.assertTrue(self.updater_mock.patched['UpdateRootfs'].called)
David Pursellf1d16a62015-03-25 13:31:04 -070080
81 def testUpdateStateful(self):
82 """Tests that update methods are called correctly."""
David Pursellf1d16a62015-03-25 13:31:04 -070083 with mock.patch('os.path.exists', return_value=True):
Bertrand SIMONNETb34a98b2015-04-22 14:30:04 -070084 flash.Flash(self.DEVICE, self.IMAGE, rootfs_update=False)
85 self.assertTrue(self.updater_mock.patched['UpdateStateful'].called)
86 self.assertFalse(self.updater_mock.patched['UpdateRootfs'].called)
David Pursellf1d16a62015-03-25 13:31:04 -070087
88 def testUpdateRootfs(self):
89 """Tests that update methods are called correctly."""
David Pursellf1d16a62015-03-25 13:31:04 -070090 with mock.patch('os.path.exists', return_value=True):
Bertrand SIMONNETb34a98b2015-04-22 14:30:04 -070091 flash.Flash(self.DEVICE, self.IMAGE, stateful_update=False)
92 self.assertFalse(self.updater_mock.patched['UpdateStateful'].called)
93 self.assertTrue(self.updater_mock.patched['UpdateRootfs'].called)
David Pursellf1d16a62015-03-25 13:31:04 -070094
95 def testMissingPayloads(self):
96 """Tests we raise FlashError when payloads are missing."""
97 with mock.patch('os.path.exists', return_value=False):
Sanika Kulkarnia8c4e3a2019-09-20 16:47:25 -070098 self.assertRaises(auto_updater_transfer.ChromiumOSTransferError,
99 flash.Flash, self.DEVICE, self.IMAGE)
David Pursellf1d16a62015-03-25 13:31:04 -0700100
Achuith Bhandarkar4d4275f2019-10-01 17:07:23 +0200101 def testFullPayload(self):
102 """Tests that we download full_payload and stateful using xBuddy."""
103 with mock.patch.object(
104 dev_server_wrapper,
105 'GetImagePathWithXbuddy',
106 return_value=('translated/xbuddy/path',
107 'resolved/xbuddy/path')) as mock_xbuddy:
108 with mock.patch('os.path.exists', return_value=True):
109 flash.Flash(self.DEVICE, self.IMAGE)
110
111 # Call to download full_payload and stateful. No other calls.
112 mock_xbuddy.assert_has_calls(
113 [mock.call('/path/to/image/full_payload', mock.ANY,
Achuith Bhandarkar12f43c72019-11-21 16:44:24 -0800114 static_dir=flash.DEVSERVER_STATIC_DIR, silent=True),
Achuith Bhandarkar4d4275f2019-10-01 17:07:23 +0200115 mock.call('/path/to/image/stateful', mock.ANY,
Achuith Bhandarkar12f43c72019-11-21 16:44:24 -0800116 static_dir=flash.DEVSERVER_STATIC_DIR, silent=True)])
Achuith Bhandarkar4d4275f2019-10-01 17:07:23 +0200117 self.assertEqual(mock_xbuddy.call_count, 2)
118
119 def testTestImage(self):
120 """Tests that we download the test image when the full payload fails."""
121 with mock.patch.object(
122 dev_server_wrapper,
123 'GetImagePathWithXbuddy',
124 side_effect=(dev_server_wrapper.ImagePathError,
125 ('translated/xbuddy/path',
126 'resolved/xbuddy/path'))) as mock_xbuddy:
127 with mock.patch('os.path.exists', return_value=True):
128 flash.Flash(self.DEVICE, self.IMAGE)
129
130 # Call to download full_payload and image. No other calls.
131 mock_xbuddy.assert_has_calls(
132 [mock.call('/path/to/image/full_payload', mock.ANY,
Achuith Bhandarkar12f43c72019-11-21 16:44:24 -0800133 static_dir=flash.DEVSERVER_STATIC_DIR, silent=True),
Achuith Bhandarkar4d4275f2019-10-01 17:07:23 +0200134 mock.call('/path/to/image', mock.ANY,
135 static_dir=flash.DEVSERVER_STATIC_DIR)])
136 self.assertEqual(mock_xbuddy.call_count, 2)
137
David Pursellf1d16a62015-03-25 13:31:04 -0700138
139class USBImagerMock(partial_mock.PartialCmdMock):
140 """Mock out USBImager."""
141 TARGET = 'chromite.cli.flash.USBImager'
142 ATTRS = ('CopyImageToDevice', 'InstallImageToDevice',
143 'ChooseRemovableDevice', 'ListAllRemovableDevices',
Bertrand SIMONNET56f773d2015-05-04 14:02:39 -0700144 'GetRemovableDeviceDescription')
David Pursellf1d16a62015-03-25 13:31:04 -0700145 VALID_IMAGE = True
146
147 def __init__(self):
148 partial_mock.PartialCmdMock.__init__(self)
149
150 def CopyImageToDevice(self, _inst, *_args, **_kwargs):
151 """Mock out CopyImageToDevice."""
152
153 def InstallImageToDevice(self, _inst, *_args, **_kwargs):
154 """Mock out InstallImageToDevice."""
155
156 def ChooseRemovableDevice(self, _inst, *_args, **_kwargs):
157 """Mock out ChooseRemovableDevice."""
158
159 def ListAllRemovableDevices(self, _inst, *_args, **_kwargs):
160 """Mock out ListAllRemovableDevices."""
161 return ['foo', 'taco', 'milk']
162
163 def GetRemovableDeviceDescription(self, _inst, *_args, **_kwargs):
164 """Mock out GetRemovableDeviceDescription."""
165
David Pursellf1d16a62015-03-25 13:31:04 -0700166
167class USBImagerTest(cros_test_lib.MockTempDirTestCase):
168 """Test the flow of flash.Flash() with USBImager."""
169 IMAGE = '/path/to/image'
170
171 def Device(self, path):
172 """Create a USB device for passing to flash.Flash()."""
173 return commandline.Device(scheme=commandline.DEVICE_SCHEME_USB,
174 path=path)
175
176 def setUp(self):
177 """Patches objects."""
178 self.usb_mock = USBImagerMock()
179 self.imager_mock = self.StartPatcher(self.usb_mock)
180 self.PatchObject(dev_server_wrapper, 'GenerateXbuddyRequest',
181 return_value='xbuddy/local/latest')
David Pursellf1d16a62015-03-25 13:31:04 -0700182 self.PatchObject(dev_server_wrapper, 'GetImagePathWithXbuddy',
Gilad Arnolde62ec902015-04-24 14:41:02 -0700183 return_value=('taco-paladin/R36/chromiumos_test_image.bin',
184 'remote/taco-paladin/R36/test'))
David Pursellf1d16a62015-03-25 13:31:04 -0700185 self.PatchObject(os.path, 'exists', return_value=True)
Bertrand SIMONNET56f773d2015-05-04 14:02:39 -0700186 self.isgpt_mock = self.PatchObject(flash, '_IsFilePathGPTDiskImage',
187 return_value=True)
David Pursellf1d16a62015-03-25 13:31:04 -0700188
189 def testLocalImagePathCopy(self):
190 """Tests that imaging methods are called correctly."""
191 with mock.patch('os.path.isfile', return_value=True):
192 flash.Flash(self.Device('/dev/foo'), self.IMAGE)
193 self.assertTrue(self.imager_mock.patched['CopyImageToDevice'].called)
194
195 def testLocalImagePathInstall(self):
196 """Tests that imaging methods are called correctly."""
197 with mock.patch('os.path.isfile', return_value=True):
198 flash.Flash(self.Device('/dev/foo'), self.IMAGE, board='taco',
199 install=True)
200 self.assertTrue(self.imager_mock.patched['InstallImageToDevice'].called)
201
202 def testLocalBadImagePath(self):
203 """Tests that using an image not having the magic bytes has prompt."""
Bertrand SIMONNET56f773d2015-05-04 14:02:39 -0700204 self.isgpt_mock.return_value = False
David Pursellf1d16a62015-03-25 13:31:04 -0700205 with mock.patch('os.path.isfile', return_value=True):
206 with mock.patch.object(cros_build_lib, 'BooleanPrompt') as mock_prompt:
207 mock_prompt.return_value = False
208 flash.Flash(self.Device('/dev/foo'), self.IMAGE)
209 self.assertTrue(mock_prompt.called)
210
211 def testNonLocalImagePath(self):
212 """Tests that we try to get the image path using xbuddy."""
213 with mock.patch.object(
214 dev_server_wrapper,
215 'GetImagePathWithXbuddy',
Gilad Arnolde62ec902015-04-24 14:41:02 -0700216 return_value=('translated/xbuddy/path',
217 'resolved/xbuddy/path')) as mock_xbuddy:
David Pursellf1d16a62015-03-25 13:31:04 -0700218 with mock.patch('os.path.isfile', return_value=False):
219 with mock.patch('os.path.isdir', return_value=False):
220 flash.Flash(self.Device('/dev/foo'), self.IMAGE)
221 self.assertTrue(mock_xbuddy.called)
222
223 def testConfirmNonRemovableDevice(self):
224 """Tests that we ask user to confirm if the device is not removable."""
225 with mock.patch.object(cros_build_lib, 'BooleanPrompt') as mock_prompt:
226 flash.Flash(self.Device('/dev/dummy'), self.IMAGE)
227 self.assertTrue(mock_prompt.called)
228
229 def testSkipPromptNonRemovableDevice(self):
230 """Tests that we skip the prompt for non-removable with --yes."""
231 with mock.patch.object(cros_build_lib, 'BooleanPrompt') as mock_prompt:
232 flash.Flash(self.Device('/dev/dummy'), self.IMAGE, yes=True)
233 self.assertFalse(mock_prompt.called)
234
235 def testChooseRemovableDevice(self):
236 """Tests that we ask user to choose a device if none is given."""
237 flash.Flash(self.Device(''), self.IMAGE)
238 self.assertTrue(self.imager_mock.patched['ChooseRemovableDevice'].called)
Bertrand SIMONNET56f773d2015-05-04 14:02:39 -0700239
240
Benjamin Gordon121a2aa2018-05-04 16:24:45 -0600241class UsbImagerOperationTest(cros_test_lib.RunCommandTestCase):
Ralph Nathan9b997232015-05-15 13:13:12 -0700242 """Tests for flash.UsbImagerOperation."""
243 # pylint: disable=protected-access
244
245 def setUp(self):
246 self.PatchObject(flash.UsbImagerOperation, '__init__', return_value=None)
247
248 def testUsbImagerOperationCalled(self):
249 """Test that flash.UsbImagerOperation is called when log level <= NOTICE."""
250 expected_cmd = ['dd', 'if=foo', 'of=bar', 'bs=4M', 'iflag=fullblock',
Frank Huang8e626432019-06-24 19:51:08 +0800251 'oflag=direct', 'conv=fdatasync']
Ralph Nathan9b997232015-05-15 13:13:12 -0700252 usb_imager = flash.USBImager('dummy_device', 'board', 'foo')
253 run_mock = self.PatchObject(flash.UsbImagerOperation, 'Run')
254 self.PatchObject(logging.Logger, 'getEffectiveLevel',
255 return_value=logging.NOTICE)
256 usb_imager.CopyImageToDevice('foo', 'bar')
257
258 # Check that flash.UsbImagerOperation.Run() is called correctly.
Mike Frysinger45602c72019-09-22 02:15:11 -0400259 run_mock.assert_called_with(cros_build_lib.sudo_run, expected_cmd,
Mike Frysinger3d5de8f2019-10-23 00:48:39 -0400260 debug_level=logging.NOTICE, encoding='utf-8',
261 update_period=0.5)
Ralph Nathan9b997232015-05-15 13:13:12 -0700262
263 def testSudoRunCommandCalled(self):
Mike Frysinger45602c72019-09-22 02:15:11 -0400264 """Test that sudo_run is called when log level > NOTICE."""
Ralph Nathan9b997232015-05-15 13:13:12 -0700265 expected_cmd = ['dd', 'if=foo', 'of=bar', 'bs=4M', 'iflag=fullblock',
Frank Huang8e626432019-06-24 19:51:08 +0800266 'oflag=direct', 'conv=fdatasync']
Ralph Nathan9b997232015-05-15 13:13:12 -0700267 usb_imager = flash.USBImager('dummy_device', 'board', 'foo')
Mike Frysinger45602c72019-09-22 02:15:11 -0400268 run_mock = self.PatchObject(cros_build_lib, 'sudo_run')
Ralph Nathan9b997232015-05-15 13:13:12 -0700269 self.PatchObject(logging.Logger, 'getEffectiveLevel',
270 return_value=logging.WARNING)
271 usb_imager.CopyImageToDevice('foo', 'bar')
272
Mike Frysinger45602c72019-09-22 02:15:11 -0400273 # Check that sudo_run() is called correctly.
Ralph Nathan9b997232015-05-15 13:13:12 -0700274 run_mock.assert_any_call(expected_cmd, debug_level=logging.NOTICE,
275 print_cmd=False)
276
277 def testPingDD(self):
278 """Test that UsbImagerOperation._PingDD() sends the correct signal."""
279 expected_cmd = ['kill', '-USR1', '5']
Mike Frysinger45602c72019-09-22 02:15:11 -0400280 run_mock = self.PatchObject(cros_build_lib, 'sudo_run')
Ralph Nathan9b997232015-05-15 13:13:12 -0700281 op = flash.UsbImagerOperation('foo')
282 op._PingDD(5)
283
Mike Frysinger45602c72019-09-22 02:15:11 -0400284 # Check that sudo_run was called correctly.
Ralph Nathan9b997232015-05-15 13:13:12 -0700285 run_mock.assert_called_with(expected_cmd, print_cmd=False)
286
287 def testGetDDPidFound(self):
288 """Check that the expected pid is returned for _GetDDPid()."""
289 expected_pid = 5
290 op = flash.UsbImagerOperation('foo')
291 self.PatchObject(osutils, 'IsChildProcess', return_value=True)
292 self.rc.AddCmdResult(partial_mock.Ignore(),
293 output='%d\n10\n' % expected_pid)
294
295 pid = op._GetDDPid()
296
297 # Check that the correct pid was returned.
298 self.assertEqual(pid, expected_pid)
299
300 def testGetDDPidNotFound(self):
301 """Check that -1 is returned for _GetDDPid() if the pids aren't valid."""
302 expected_pid = -1
303 op = flash.UsbImagerOperation('foo')
304 self.PatchObject(osutils, 'IsChildProcess', return_value=False)
305 self.rc.AddCmdResult(partial_mock.Ignore(), output='5\n10\n')
306
307 pid = op._GetDDPid()
308
309 # Check that the correct pid was returned.
310 self.assertEqual(pid, expected_pid)
311
312
Bertrand SIMONNET56f773d2015-05-04 14:02:39 -0700313class FlashUtilTest(cros_test_lib.MockTempDirTestCase):
314 """Tests the helpers from cli.flash."""
315
316 def testChooseImage(self):
317 """Tests that we can detect a GPT image."""
318 # pylint: disable=protected-access
319
320 with self.PatchObject(flash, '_IsFilePathGPTDiskImage', return_value=True):
321 # No images defined. Choosing the image should raise an error.
322 with self.assertRaises(ValueError):
323 flash._ChooseImageFromDirectory(self.tempdir)
324
325 file_a = os.path.join(self.tempdir, 'a')
326 osutils.Touch(file_a)
327 # Only one image available, it should be selected automatically.
328 self.assertEqual(file_a, flash._ChooseImageFromDirectory(self.tempdir))
329
330 osutils.Touch(os.path.join(self.tempdir, 'b'))
331 file_c = os.path.join(self.tempdir, 'c')
332 osutils.Touch(file_c)
333 osutils.Touch(os.path.join(self.tempdir, 'd'))
334
335 # Multiple images available, we should ask the user to select the right
336 # image.
337 with self.PatchObject(cros_build_lib, 'GetChoice', return_value=2):
338 self.assertEqual(file_c, flash._ChooseImageFromDirectory(self.tempdir))
Mike Frysinger32759e42016-12-21 18:40:16 -0500339
340 def testIsFilePathGPTDiskImage(self):
341 """Tests the GPT image probing."""
342 # pylint: disable=protected-access
343
Mike Frysinger3d5de8f2019-10-23 00:48:39 -0400344 INVALID_PMBR = b' ' * 0x200
345 INVALID_GPT = b' ' * 0x200
346 VALID_PMBR = (b' ' * 0x1fe) + b'\x55\xaa'
347 VALID_GPT = b'EFI PART' + (b' ' * 0x1f8)
Mike Frysinger32759e42016-12-21 18:40:16 -0500348 TESTCASES = (
349 (False, False, INVALID_PMBR + INVALID_GPT),
350 (False, False, VALID_PMBR + INVALID_GPT),
351 (False, True, INVALID_PMBR + VALID_GPT),
352 (True, True, VALID_PMBR + VALID_GPT),
353 )
354
355 img = os.path.join(self.tempdir, 'img.bin')
356 for exp_pmbr_t, exp_pmbr_f, data in TESTCASES:
Mike Frysinger3d5de8f2019-10-23 00:48:39 -0400357 osutils.WriteFile(img, data, mode='wb')
Mike Frysinger32759e42016-12-21 18:40:16 -0500358 self.assertEqual(
359 flash._IsFilePathGPTDiskImage(img, require_pmbr=True), exp_pmbr_t)
360 self.assertEqual(
361 flash._IsFilePathGPTDiskImage(img, require_pmbr=False), exp_pmbr_f)