1#!/usr/bin/env vpython3 2# Copyright 2022 The Chromium Authors 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5"""File for testing flash_device.py.""" 6 7import os 8import unittest 9import unittest.mock as mock 10 11import boot_device 12import flash_device 13 14_TEST_IMAGE_DIR = 'test/image/dir' 15_TEST_PRODUCT = 'test_product' 16_TEST_VERSION = 'test.version' 17 18 19# pylint: disable=too-many-public-methods,protected-access 20class FlashDeviceTest(unittest.TestCase): 21 """Unittests for flash_device.py.""" 22 23 def setUp(self) -> None: 24 context_mock = mock.Mock() 25 context_mock.__enter__ = mock.Mock(return_value=None) 26 context_mock.__exit__ = mock.Mock(return_value=None) 27 ffx_mock = mock.Mock() 28 ffx_mock.returncode = 0 29 ffx_patcher = mock.patch('common.run_ffx_command', 30 return_value=ffx_mock) 31 sdk_hash_patcher = mock.patch('flash_device.get_sdk_hash', 32 return_value=(_TEST_PRODUCT, 33 _TEST_VERSION)) 34 swarming_patcher = mock.patch('flash_device.running_unattended', 35 return_value=False) 36 time_sleep = mock.patch('time.sleep') 37 self._ffx_mock = ffx_patcher.start() 38 self._sdk_hash_mock = sdk_hash_patcher.start() 39 self._swarming_mock = swarming_patcher.start() 40 self._time_sleep = time_sleep.start() 41 self.addCleanup(self._ffx_mock.stop) 42 self.addCleanup(self._sdk_hash_mock.stop) 43 self.addCleanup(self._swarming_mock.stop) 44 self.addCleanup(self._time_sleep.stop) 45 46 def test_update_required_on_ignore_returns_immediately(self) -> None: 47 """Test |os_check|='ignore' skips all checks.""" 48 result, new_image_dir = flash_device._update_required( 49 'ignore', 'some-image-dir', None) 50 51 self.assertFalse(result) 52 self.assertEqual(new_image_dir, 'some-image-dir') 53 54 def test_update_required_raises_value_error_if_no_image_dir(self) -> None: 55 """Test |os_check|!='ignore' checks that image dir is non-Falsey.""" 56 with self.assertRaises(ValueError): 57 flash_device._update_required('update', None, None) 58 59 def test_update_required_logs_missing_image_dir(self) -> None: 60 """Test |os_check|!='ignore' warns if image dir does not exist.""" 61 with mock.patch('os.path.exists', return_value=False), \ 62 mock.patch('flash_device.find_image_in_sdk'), \ 63 mock.patch('flash_device._get_system_info'), \ 64 self.assertLogs() as logger: 65 flash_device._update_required('update', 'some/image/dir', None) 66 self.assertIn('image directory does not exist', logger.output[0]) 67 68 def test_update_required_searches_and_returns_sdk_if_image_found(self 69 ) -> None: 70 """Test |os_check|!='ignore' searches for image dir in SDK.""" 71 with mock.patch('os.path.exists', return_value=False), \ 72 mock.patch('flash_device.find_image_in_sdk') as mock_find, \ 73 mock.patch('flash_device._get_system_info'), \ 74 mock.patch('common.SDK_ROOT', 'path/to/sdk/dir'), \ 75 self.assertLogs(): 76 mock_find.return_value = 'path/to/image/dir' 77 update_required, new_image_dir = flash_device._update_required( 78 'update', 'product-bundle', None, None) 79 self.assertTrue(update_required) 80 self.assertEqual(new_image_dir, 'path/to/image/dir') 81 mock_find.assert_called_once_with('product-bundle') 82 83 def test_update_required_raises_file_not_found_error(self) -> None: 84 """Test |os_check|!='ignore' raises FileNotFoundError if no path.""" 85 with mock.patch('os.path.exists', return_value=False), \ 86 mock.patch('flash_device.find_image_in_sdk', 87 return_value=None), \ 88 mock.patch('common.SDK_ROOT', 'path/to/sdk/dir'), \ 89 self.assertLogs(), \ 90 self.assertRaises(FileNotFoundError): 91 flash_device._update_required('update', 'product-bundle', None) 92 93 def test_update_ignore(self) -> None: 94 """Test setting |os_check| to 'ignore'.""" 95 96 flash_device.update(_TEST_IMAGE_DIR, 'ignore', None) 97 self.assertEqual(self._ffx_mock.call_count, 0) 98 self.assertEqual(self._sdk_hash_mock.call_count, 0) 99 100 def test_dir_unspecified_value_error(self) -> None: 101 """Test ValueError raised when system_image_dir unspecified.""" 102 103 with self.assertRaises(ValueError): 104 flash_device.update(None, 'check', None) 105 106 def test_update_system_info_match(self) -> None: 107 """Test no update when |os_check| is 'check' and system info matches.""" 108 109 with mock.patch('os.path.exists', return_value=True): 110 self._ffx_mock.return_value.stdout = \ 111 '{"build": {"version": "%s", ' \ 112 '"product": "%s"}}' % (_TEST_VERSION, _TEST_PRODUCT) 113 flash_device.update(_TEST_IMAGE_DIR, 'check', None) 114 self.assertEqual(self._ffx_mock.call_count, 1) 115 self.assertEqual(self._sdk_hash_mock.call_count, 1) 116 117 def test_update_system_info_catches_boot_failure(self) -> None: 118 """Test update when |os_check=check| catches boot_device exceptions.""" 119 120 self._swarming_mock.return_value = True 121 with mock.patch('os.path.exists', return_value=True), \ 122 mock.patch('flash_device.boot_device') as mock_boot, \ 123 mock.patch('flash_device.get_system_info') as mock_sys_info, \ 124 mock.patch('flash_device.subprocess.run'): 125 mock_boot.side_effect = boot_device.StateTransitionError( 126 'Incorrect state') 127 self._ffx_mock.return_value.stdout = \ 128 '{"build": {"version": "wrong.version", ' \ 129 '"product": "wrong.product"}}' 130 flash_device.update(_TEST_IMAGE_DIR, 'check', None) 131 mock_boot.assert_called_with(mock.ANY, 132 boot_device.BootMode.REGULAR, None) 133 self.assertEqual(self._ffx_mock.call_count, 1) 134 135 # get_system_info should not even be called due to early exit. 136 mock_sys_info.assert_not_called() 137 138 def test_update_system_info_mismatch(self) -> None: 139 """Test update when |os_check| is 'check' and system info does not 140 match.""" 141 142 self._swarming_mock.return_value = True 143 with mock.patch('os.path.exists', return_value=True), \ 144 mock.patch('flash_device.boot_device') as mock_boot, \ 145 mock.patch('flash_device.subprocess.run'): 146 self._ffx_mock.return_value.stdout = \ 147 '{"build": {"version": "wrong.version", ' \ 148 '"product": "wrong.product"}}' 149 flash_device.update(_TEST_IMAGE_DIR, 'check', None) 150 mock_boot.assert_called_with(mock.ANY, 151 boot_device.BootMode.REGULAR, None) 152 self.assertEqual(self._ffx_mock.call_count, 2) 153 154 def test_incorrect_target_info(self) -> None: 155 """Test update when |os_check| is 'check' and system info was not 156 retrieved.""" 157 with mock.patch('os.path.exists', return_value=True): 158 self._ffx_mock.return_value.stdout = '{"unexpected": "badtitle"}' 159 flash_device.update(_TEST_IMAGE_DIR, 'check', None) 160 self.assertEqual(self._ffx_mock.call_count, 2) 161 162 def test_update_with_serial_num(self) -> None: 163 """Test update when |serial_num| is specified.""" 164 165 with mock.patch('time.sleep'), \ 166 mock.patch('os.path.exists', return_value=True), \ 167 mock.patch('flash_device.boot_device') as mock_boot: 168 flash_device.update(_TEST_IMAGE_DIR, 'update', None, 'test_serial') 169 mock_boot.assert_called_with(mock.ANY, 170 boot_device.BootMode.BOOTLOADER, 171 'test_serial') 172 self.assertEqual(self._ffx_mock.call_count, 1) 173 174 def test_reboot_failure(self) -> None: 175 """Test update when |serial_num| is specified.""" 176 self._ffx_mock.return_value.returncode = 1 177 with mock.patch('time.sleep'), \ 178 mock.patch('os.path.exists', return_value=True), \ 179 mock.patch('flash_device.running_unattended', 180 return_value=True), \ 181 mock.patch('flash_device.boot_device'): 182 required, _ = flash_device._update_required( 183 'check', _TEST_IMAGE_DIR, None) 184 self.assertEqual(required, True) 185 186 def test_update_on_swarming(self) -> None: 187 """Test update on swarming bots.""" 188 189 self._swarming_mock.return_value = True 190 with mock.patch('time.sleep'), \ 191 mock.patch('os.path.exists', return_value=True), \ 192 mock.patch('flash_device.boot_device') as mock_boot, \ 193 mock.patch('subprocess.run'): 194 flash_device.update(_TEST_IMAGE_DIR, 'update', None, 'test_serial') 195 mock_boot.assert_called_with(mock.ANY, 196 boot_device.BootMode.BOOTLOADER, 197 'test_serial') 198 self.assertEqual(self._ffx_mock.call_count, 1) 199 200 def test_main(self) -> None: 201 """Tests |main| function.""" 202 203 with mock.patch('sys.argv', 204 ['flash_device.py', '--os-check', 'ignore']): 205 with mock.patch.dict(os.environ, {}): 206 flash_device.main() 207 self.assertEqual(self._ffx_mock.call_count, 0) 208# pylint: enable=too-many-public-methods,protected-access 209 210 211if __name__ == '__main__': 212 unittest.main() 213