| # Copyright 2017 The Android Open Source Project |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| """Unit test for atft manager.""" |
| import base64 |
| import unittest |
| |
| import atftman |
| |
| from atftman import EncryptionAlgorithm |
| from atftman import ProductInfo |
| from atftman import ProvisionState |
| from atftman import ProvisionStatus |
| from fastboot_exceptions import DeviceNotFoundException |
| from fastboot_exceptions import FastbootFailure |
| from fastboot_exceptions import NoAlgorithmAvailableException |
| from fastboot_exceptions import ProductAttributesFileFormatError |
| from fastboot_exceptions import ProductNotSpecifiedException |
| from mock import call |
| from mock import MagicMock |
| from mock import patch |
| |
| files = [] |
| |
| |
| class AtftManTest(unittest.TestCase): |
| ATFA_TEST_SERIAL = 'ATFA_TEST_SERIAL' |
| TEST_TMP_FOLDER = '/tmp/TMPTEST/' |
| TEST_SERIAL = 'TEST_SERIAL' |
| TEST_SERIAL2 = 'TEST_SERIAL2' |
| TEST_SERIAL3 = 'TEST_SERIAL3' |
| TEST_UUID = 'TEST-UUID' |
| TEST_LOCATION = 'BUS1-PORT1' |
| TEST_LOCATION2 = 'BUS2-PORT1' |
| TEST_LOCATION3 = 'BUS1-PORT2' |
| TEST_ID = '00000000000000000000000000000000' |
| TEST_ID_ARRAY = bytearray.fromhex(TEST_ID) |
| TEST_NAME = 'name' |
| TEST_FILE_NAME = 'filename' |
| TEST_ATTRIBUTE_ARRAY = bytearray(1052) |
| TEST_VBOOT_KEY_ARRAY = bytearray(128) |
| TEST_ATTRIBUTE_STRING = base64.standard_b64encode(TEST_ATTRIBUTE_ARRAY) |
| TEST_VBOOT_KEY_STRING = base64.standard_b64encode(TEST_VBOOT_KEY_ARRAY) |
| |
| class FastbootDeviceTemplate(object): |
| |
| @staticmethod |
| def ListDevices(): |
| pass |
| |
| def __init__(self, serial_number): |
| self.serial_number = serial_number |
| |
| def Oem(self, oem_command, err_to_out): |
| pass |
| |
| def Upload(self, file_path): |
| pass |
| |
| def GetVar(self, var): |
| pass |
| |
| def Download(self, file_path): |
| pass |
| |
| def Disconnect(self): |
| pass |
| |
| def GetHostOs(self): |
| return 'Windows' |
| |
| def __del__(self): |
| pass |
| |
| def MockInit(self, serial_number): |
| mock_instance = MagicMock() |
| mock_instance.serial_number = serial_number |
| mock_instance.GetHostOs = MagicMock() |
| mock_instance.GetHostOs.return_value = 'Windows' |
| return mock_instance |
| |
| def setUp(self): |
| self.mock_serial_mapper = MagicMock() |
| self.mock_serial_instance = MagicMock() |
| self.mock_serial_mapper.return_value = self.mock_serial_instance |
| self.mock_serial_instance.get_serial_map.return_value = [] |
| self.status_map = {} |
| self.mock_timer_instance = None |
| self.configs = {} |
| self.configs['ATFA_REBOOT_TIMEOUT'] = 30 |
| self.configs['DEFAULT_KEY_THRESHOLD'] = 100 |
| |
| # Test AtftManager.ListDevices |
| class MockInstantTimer(object): |
| def __init__(self, timeout, callback): |
| self.timeout = timeout |
| self.callback = callback |
| |
| def start(self): |
| self.callback() |
| |
| def MockCreateInstantTimer(self, timeout, callback): |
| return self.MockInstantTimer(timeout, callback) |
| |
| @patch('threading.Timer') |
| @patch('__main__.AtftManTest.FastbootDeviceTemplate.ListDevices') |
| def testListDevicesNormal(self, mock_list_devices, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._SetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_list_devices.return_value = [self.TEST_SERIAL, self.ATFA_TEST_SERIAL] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev.serial_number, self.ATFA_TEST_SERIAL) |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| self.assertEqual(atft_manager.target_devs[0].serial_number, |
| self.TEST_SERIAL) |
| self.assertEqual(None, atft_manager._atfa_dev_setting) |
| self.assertEqual(True, atft_manager._atfa_reboot_lock.acquire(False)) |
| atft_manager._SetOs.assert_not_called() |
| |
| @patch('threading.Timer') |
| @patch('__main__.AtftManTest.FastbootDeviceTemplate.ListDevices') |
| def testListDevicesNormalSetOs(self, mock_list_devices, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._SetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Linux' |
| mock_list_devices.return_value = [self.TEST_SERIAL, self.ATFA_TEST_SERIAL] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| atft_manager._GetOs.assert_called() |
| atft_manager._SetOs.assert_called_once_with( |
| atft_manager.atfa_dev, 'Windows') |
| self.assertEqual(atft_manager.atfa_dev.serial_number, self.ATFA_TEST_SERIAL) |
| |
| @patch('threading.Timer') |
| @patch('__main__.AtftManTest.FastbootDeviceTemplate.ListDevices') |
| def testListDevicesATFA(self, mock_list_devices, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_list_devices.return_value = [self.ATFA_TEST_SERIAL] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev.serial_number, self.ATFA_TEST_SERIAL) |
| self.assertEqual(0, len(atft_manager.target_devs)) |
| |
| @patch('threading.Timer') |
| @patch('__main__.AtftManTest.FastbootDeviceTemplate.ListDevices') |
| def testListDevicesTarget(self, mock_list_devices, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_list_devices.return_value = [self.TEST_SERIAL] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev, None) |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| self.assertEqual(atft_manager.target_devs[0].serial_number, |
| self.TEST_SERIAL) |
| |
| @patch('threading.Timer') |
| @patch('__main__.AtftManTest.FastbootDeviceTemplate.ListDevices') |
| def testListDevicesMultipleTargets(self, mock_list_devices, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_list_devices.return_value = [self.TEST_SERIAL, self.TEST_SERIAL2] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev, None) |
| self.assertEqual(2, len(atft_manager.target_devs)) |
| self.assertEqual(atft_manager.target_devs[0].serial_number, |
| self.TEST_SERIAL) |
| self.assertEqual(atft_manager.target_devs[1].serial_number, |
| self.TEST_SERIAL2) |
| |
| @patch('threading.Timer') |
| def testListDevicesChangeNorm(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [self.ATFA_TEST_SERIAL] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| mock_fastboot.ListDevices.return_value = [self.TEST_SERIAL] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev, None) |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| self.assertEqual(atft_manager.target_devs[0].serial_number, |
| self.TEST_SERIAL) |
| self.assertEqual(2, mock_fastboot.call_count) |
| |
| @patch('threading.Timer') |
| def testListDevicesChangeAdd(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [self.ATFA_TEST_SERIAL] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| mock_fastboot.ListDevices.return_value = [ |
| self.ATFA_TEST_SERIAL, self.TEST_SERIAL |
| ] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev.serial_number, self.ATFA_TEST_SERIAL) |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| self.assertEqual(atft_manager.target_devs[0].serial_number, |
| self.TEST_SERIAL) |
| self.assertEqual(2, mock_fastboot.call_count) |
| |
| @patch('threading.Timer') |
| def testListDevicesChangeAddATFA(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [self.TEST_SERIAL] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| mock_fastboot.ListDevices.return_value = [ |
| self.ATFA_TEST_SERIAL, self.TEST_SERIAL |
| ] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev.serial_number, self.ATFA_TEST_SERIAL) |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| self.assertEqual(atft_manager.target_devs[0].serial_number, |
| self.TEST_SERIAL) |
| self.assertEqual(2, mock_fastboot.call_count) |
| |
| @patch('threading.Timer') |
| def testListDevicesChangeCommon(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [ |
| self.ATFA_TEST_SERIAL, self.TEST_SERIAL |
| ] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| mock_fastboot.ListDevices.return_value = [ |
| self.TEST_SERIAL2, self.TEST_SERIAL |
| ] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev, None) |
| self.assertEqual(2, len(atft_manager.target_devs)) |
| self.assertEqual(atft_manager.target_devs[0].serial_number, |
| self.TEST_SERIAL) |
| self.assertEqual(atft_manager.target_devs[1].serial_number, |
| self.TEST_SERIAL2) |
| self.assertEqual(3, mock_fastboot.call_count) |
| |
| @patch('threading.Timer') |
| def testListDevicesChangeCommonATFA(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [ |
| self.ATFA_TEST_SERIAL, self.TEST_SERIAL |
| ] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| mock_fastboot.ListDevices.return_value = [ |
| self.ATFA_TEST_SERIAL, self.TEST_SERIAL2 |
| ] |
| atft_manager.ListDevices() |
| # First refresh, TEST_SERIAL should disappear |
| self.assertEqual(0, len(atft_manager.target_devs)) |
| # Second refresh, TEST_SERIAL2 should be added |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev.serial_number, self.ATFA_TEST_SERIAL) |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| self.assertEqual(atft_manager.target_devs[0].serial_number, |
| self.TEST_SERIAL2) |
| self.assertEqual(3, mock_fastboot.call_count) |
| |
| @patch('threading.Timer') |
| def testListDevicesRemoveATFA(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [ |
| self.ATFA_TEST_SERIAL, self.TEST_SERIAL |
| ] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| mock_fastboot.ListDevices.return_value = [self.TEST_SERIAL] |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev, None) |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| self.assertEqual(atft_manager.target_devs[0].serial_number, |
| self.TEST_SERIAL) |
| |
| @patch('threading.Timer') |
| def testListDevicesRemoveDevice(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [ |
| self.ATFA_TEST_SERIAL, self.TEST_SERIAL |
| ] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| mock_fastboot.ListDevices.return_value = [self.ATFA_TEST_SERIAL] |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev.serial_number, self.ATFA_TEST_SERIAL) |
| self.assertEqual(0, len(atft_manager.target_devs)) |
| |
| @patch('threading.Timer') |
| def testListDevicesPendingRemove(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [self.TEST_SERIAL] |
| atft_manager.ListDevices() |
| # Just appear once, should not be in target device list. |
| self.assertEqual(0, len(atft_manager.target_devs)) |
| |
| @patch('threading.Timer') |
| def testListDevicesPendingAdd(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [self.TEST_SERIAL] |
| atft_manager.ListDevices() |
| self.assertEqual(0, len(atft_manager.target_devs)) |
| mock_fastboot.ListDevices.return_value = [ |
| self.TEST_SERIAL, self.TEST_SERIAL2 |
| ] |
| # TEST_SERIAL appears twice, should be in the list. |
| # TEST_SERIAL2 just appears once, should not be in the list. |
| atft_manager.ListDevices() |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| |
| @patch('threading.Timer') |
| def testListDevicesPendingTemp(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager(mock_fastboot, self.mock_serial_mapper, |
| self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [self.TEST_SERIAL] |
| atft_manager.ListDevices() |
| self.assertEqual(0, len(atft_manager.target_devs)) |
| mock_fastboot.ListDevices.return_value = [self.TEST_SERIAL2] |
| atft_manager.ListDevices() |
| # Nothing appears twice. |
| self.assertEqual(0, len(atft_manager.target_devs)) |
| |
| def mockSetSerialMapper(self, serial_map): |
| self.serial_map = {} |
| for serial in serial_map: |
| self.serial_map[serial.lower()] = serial_map[serial] |
| |
| def mockGetLocation(self, serial): |
| serial_lower = serial.lower() |
| if serial_lower in self.serial_map: |
| return self.serial_map[serial_lower] |
| return None |
| |
| @patch('threading.Timer') |
| def testListDevicesLocation(self, mock_create_timer): |
| mock_create_timer.side_effect = self.MockCreateInstantTimer |
| mock_serial_mapper = MagicMock() |
| smap = { |
| self.ATFA_TEST_SERIAL: self.TEST_LOCATION, |
| self.TEST_SERIAL: self.TEST_LOCATION2 |
| } |
| mock_serial_instance = MagicMock() |
| mock_serial_mapper.return_value = mock_serial_instance |
| mock_serial_instance.refresh_serial_map.side_effect = ( |
| lambda serial_map=smap: self.mockSetSerialMapper(serial_map)) |
| mock_serial_instance.get_location.side_effect = self.mockGetLocation |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager( |
| mock_fastboot, mock_serial_mapper, self.configs) |
| atft_manager._GetOs = MagicMock() |
| atft_manager._GetOs.return_value = 'Windows' |
| mock_fastboot.ListDevices.return_value = [ |
| self.ATFA_TEST_SERIAL, self.TEST_SERIAL |
| ] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(atft_manager.atfa_dev.location, self.TEST_LOCATION) |
| self.assertEqual(atft_manager.target_devs[0].location, self.TEST_LOCATION2) |
| |
| # Test _SortTargetDevices |
| def testSortDevicesDefault(self): |
| mock_serial_mapper = MagicMock() |
| mock_serial_instance = MagicMock() |
| mock_serial_mapper.return_value = mock_serial_instance |
| smap = { |
| self.TEST_SERIAL: self.TEST_LOCATION, |
| self.TEST_SERIAL2: self.TEST_LOCATION2, |
| self.TEST_SERIAL3: self.TEST_LOCATION3 |
| } |
| mock_serial_instance.refresh_serial_map.side_effect = ( |
| lambda serial_map=smap: self.mockSetSerialMapper(serial_map)) |
| mock_serial_instance.get_location.side_effect = self.mockGetLocation |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager( |
| mock_fastboot, mock_serial_mapper, self.configs) |
| mock_fastboot.ListDevices.return_value = [ |
| self.TEST_SERIAL, self.TEST_SERIAL2, self.TEST_SERIAL3 |
| ] |
| atft_manager.ListDevices() |
| atft_manager.ListDevices() |
| self.assertEqual(3, len(atft_manager.target_devs)) |
| self.assertEqual(self.TEST_LOCATION, atft_manager.target_devs[0].location) |
| self.assertEqual(self.TEST_LOCATION3, atft_manager.target_devs[1].location) |
| self.assertEqual(self.TEST_LOCATION2, atft_manager.target_devs[2].location) |
| |
| def testSortDevicesLocation(self): |
| mock_serial_mapper = MagicMock() |
| mock_serial_instance = MagicMock() |
| mock_serial_mapper.return_value = mock_serial_instance |
| smap = { |
| self.TEST_SERIAL: self.TEST_LOCATION, |
| self.TEST_SERIAL2: self.TEST_LOCATION2, |
| self.TEST_SERIAL3: self.TEST_LOCATION3 |
| } |
| mock_serial_instance.refresh_serial_map.side_effect = ( |
| lambda serial_map=smap: self.mockSetSerialMapper(serial_map)) |
| mock_serial_instance.get_location.side_effect = self.mockGetLocation |
| mock_fastboot = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| atft_manager = atftman.AtftManager( |
| mock_fastboot, mock_serial_mapper, self.configs) |
| mock_fastboot.ListDevices.return_value = [ |
| self.TEST_SERIAL, self.TEST_SERIAL2, self.TEST_SERIAL3 |
| ] |
| atft_manager.ListDevices(atft_manager.SORT_BY_LOCATION) |
| atft_manager.ListDevices(atft_manager.SORT_BY_LOCATION) |
| self.assertEqual(3, len(atft_manager.target_devs)) |
| self.assertEqual(self.TEST_LOCATION, atft_manager.target_devs[0].location) |
| self.assertEqual(self.TEST_LOCATION3, atft_manager.target_devs[1].location) |
| self.assertEqual(self.TEST_LOCATION2, atft_manager.target_devs[2].location) |
| |
| def testSortDevicesSerial(self): |
| mock_serial_mapper = MagicMock() |
| mock_serial_instance = MagicMock() |
| mock_serial_mapper.return_value = mock_serial_instance |
| smap = { |
| self.TEST_SERIAL: self.TEST_LOCATION, |
| self.TEST_SERIAL2: self.TEST_LOCATION2, |
| self.TEST_SERIAL3: self.TEST_LOCATION3 |
| } |
| mock_serial_instance.refresh_serial_map.side_effect = ( |
| lambda serial_map=smap: self.mockSetSerialMapper(serial_map)) |
| mock_serial_instance.get_location.side_effect = self.mockGetLocation |
| mock_fastboot = MagicMock() |
| mock_fastboot_instance = MagicMock() |
| mock_fastboot.side_effect = self.MockInit |
| mock_fastboot.return_value = mock_fastboot_instance |
| mock_fastboot_instance.GetVar = MagicMock() |
| atft_manager = atftman.AtftManager( |
| mock_fastboot, mock_serial_mapper, self.configs) |
| mock_fastboot.ListDevices.return_value = [ |
| self.TEST_SERIAL2, self.TEST_SERIAL3, self.TEST_SERIAL |
| ] |
| atft_manager.ListDevices(atft_manager.SORT_BY_SERIAL) |
| atft_manager.ListDevices(atft_manager.SORT_BY_SERIAL) |
| self.assertEqual(3, len(atft_manager.target_devs)) |
| self.assertEqual(self.TEST_SERIAL, |
| atft_manager.target_devs[0].serial_number) |
| self.assertEqual(self.TEST_SERIAL2, |
| atft_manager.target_devs[1].serial_number) |
| self.assertEqual(self.TEST_SERIAL3, |
| atft_manager.target_devs[2].serial_number) |
| |
| # Test AtftManager.TransferContent |
| |
| @staticmethod |
| def _AppendFile(file_path): |
| files.append(file_path) |
| |
| @staticmethod |
| def _CheckFile(file_path): |
| assert file_path in files |
| return True |
| |
| @staticmethod |
| def _RemoveFile(file_path): |
| assert file_path in files |
| files.remove(file_path) |
| |
| @patch('os.rmdir') |
| @patch('os.remove') |
| @patch('os.path.exists') |
| @patch('tempfile.mkdtemp') |
| @patch('uuid.uuid1') |
| @patch('__main__.AtftManTest.FastbootDeviceTemplate.Upload') |
| @patch('__main__.AtftManTest.FastbootDeviceTemplate.Download') |
| def testTransferContentNormal(self, mock_download, mock_upload, mock_uuid, |
| mock_create_folder, mock_exists, mock_remove, |
| mock_rmdir): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| # upload (to fs): create a temporary file |
| mock_upload.side_effect = AtftManTest._AppendFile |
| # download (from fs): check if the temporary file exists |
| mock_download.side_effect = AtftManTest._CheckFile |
| mock_exists.side_effect = AtftManTest._CheckFile |
| # remove: remove the file |
| mock_remove.side_effect = AtftManTest._RemoveFile |
| mock_rmdir.side_effect = AtftManTest._RemoveFile |
| mock_create_folder.return_value = self.TEST_TMP_FOLDER |
| files.append(self.TEST_TMP_FOLDER) |
| mock_uuid.return_value = self.TEST_UUID |
| tmp_path = self.TEST_TMP_FOLDER + self.TEST_UUID |
| src = self.FastbootDeviceTemplate(self.TEST_SERIAL) |
| dst = self.FastbootDeviceTemplate(self.TEST_SERIAL) |
| atft_manager.TransferContent(src, dst) |
| src.Upload.assert_called_once_with(tmp_path) |
| src.Download.assert_called_once_with(tmp_path) |
| # we should have no temporary file at the end |
| self.assertTrue(not files) |
| |
| # Test AtftManager._ChooseAlgorithm |
| def testChooseAlgorithm(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| p256 = atftman.EncryptionAlgorithm.ALGORITHM_P256 |
| curve = atftman.EncryptionAlgorithm.ALGORITHM_CURVE25519 |
| algorithm = atft_manager._ChooseAlgorithm([p256, curve]) |
| self.assertEqual(curve, algorithm) |
| |
| def testChooseAlgorithmP256(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| p256 = atftman.EncryptionAlgorithm.ALGORITHM_P256 |
| algorithm = atft_manager._ChooseAlgorithm([p256]) |
| self.assertEqual(p256, algorithm) |
| |
| def testChooseAlgorithmCurve(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| curve = atftman.EncryptionAlgorithm.ALGORITHM_CURVE25519 |
| algorithm = atft_manager._ChooseAlgorithm([curve]) |
| self.assertEqual(curve, algorithm) |
| |
| def testChooseAlgorithmException(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| with self.assertRaises(NoAlgorithmAvailableException): |
| atft_manager._ChooseAlgorithm([]) |
| |
| def testChooseAlgorithmExceptionNoAvailable(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| with self.assertRaises(NoAlgorithmAvailableException): |
| atft_manager._ChooseAlgorithm(['abcd']) |
| |
| # Test AtftManager._GetAlgorithmList |
| def testGetAlgorithmList(self): |
| mock_target = MagicMock() |
| mock_target.GetVar.return_value = '1:p256,2:curve25519' |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| algorithm_list = atft_manager._GetAlgorithmList(mock_target) |
| self.assertEqual(2, len(algorithm_list)) |
| self.assertEqual(1, algorithm_list[0]) |
| self.assertEqual(2, algorithm_list[1]) |
| |
| # Test DeviceInfo.__eq__ |
| def testDeviceInfoEqual(self): |
| test_device1 = atftman.DeviceInfo(None, self.TEST_SERIAL, |
| self.TEST_LOCATION) |
| test_device2 = atftman.DeviceInfo(None, self.TEST_SERIAL2, |
| self.TEST_LOCATION2) |
| test_device3 = atftman.DeviceInfo(None, self.TEST_SERIAL, |
| self.TEST_LOCATION2) |
| test_device4 = atftman.DeviceInfo(None, self.TEST_SERIAL2, |
| self.TEST_LOCATION) |
| test_device5 = atftman.DeviceInfo(None, self.TEST_SERIAL, |
| self.TEST_LOCATION) |
| self.assertEqual(test_device1, test_device5) |
| self.assertNotEqual(test_device1, test_device2) |
| self.assertNotEqual(test_device1, test_device3) |
| self.assertNotEqual(test_device1, test_device4) |
| |
| # Test DeviceInfo.Copy |
| def testDeviceInfoCopy(self): |
| test_device1 = atftman.DeviceInfo(None, self.TEST_SERIAL, |
| self.TEST_LOCATION) |
| test_device2 = atftman.DeviceInfo(None, self.TEST_SERIAL2, |
| self.TEST_LOCATION2) |
| test_device3 = test_device1.Copy() |
| self.assertEqual(test_device3, test_device1) |
| self.assertNotEqual(test_device3, test_device2) |
| |
| # Test AtfaDeviceManager.CheckStatus |
| def testCheckStatus(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_atfa_dev = MagicMock() |
| atft_manager.atfa_dev = mock_atfa_dev |
| test_atfa_device_manager = atftman.AtfaDeviceManager(atft_manager) |
| atft_manager.product_info = MagicMock() |
| atft_manager.product_info.product_id = self.TEST_ID |
| mock_atfa_dev.Oem.return_value = 'TEST\n(bootloader) 100\nTEST' |
| test_number = test_atfa_device_manager.CheckStatus() |
| mock_atfa_dev.Oem.assert_called_once_with('num-keys ' + self.TEST_ID, True) |
| self.assertEqual(100, atft_manager.GetATFAKeysLeft()) |
| |
| def testCheckStatusCRLF(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_atfa_dev = MagicMock() |
| atft_manager.atfa_dev = mock_atfa_dev |
| test_atfa_device_manager = atftman.AtfaDeviceManager(atft_manager) |
| atft_manager.product_info = MagicMock() |
| atft_manager.product_info.product_id = self.TEST_ID |
| mock_atfa_dev.Oem.return_value = 'TEST\r\n(bootloader) 100\r\nTEST' |
| test_number = test_atfa_device_manager.CheckStatus() |
| mock_atfa_dev.Oem.assert_called_once_with('num-keys ' + self.TEST_ID, True) |
| self.assertEqual(100, atft_manager.GetATFAKeysLeft()) |
| |
| def testCheckStatusNoProductId(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_atfa_dev = MagicMock() |
| atft_manager.atfa_dev = mock_atfa_dev |
| atft_manager.product_info = None |
| test_atfa_device_manager = atftman.AtfaDeviceManager(atft_manager) |
| mock_atfa_dev.Oem.return_value = 'TEST\r\n(bootloader) 100\r\nTEST' |
| with self.assertRaises(ProductNotSpecifiedException): |
| test_atfa_device_manager.CheckStatus() |
| |
| def testCheckStatusNoATFA(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager.atfa_dev = None |
| atft_manager.product_info = MagicMock() |
| atft_manager.product_info.product_id = self.TEST_ID |
| test_atfa_device_manager = atftman.AtfaDeviceManager(atft_manager) |
| with self.assertRaises(DeviceNotFoundException): |
| test_atfa_device_manager.CheckStatus() |
| |
| def testCheckStatusInvalidFormat(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_atfa_dev = MagicMock() |
| atft_manager.atfa_dev = mock_atfa_dev |
| test_atfa_device_manager = atftman.AtfaDeviceManager(atft_manager) |
| atft_manager.product_info = MagicMock() |
| atft_manager.product_info.product_id = self.TEST_ID |
| mock_atfa_dev.Oem.return_value = 'TEST\nTEST' |
| with self.assertRaises(FastbootFailure): |
| test_atfa_device_manager.CheckStatus() |
| |
| def testCheckStatusInvalidNumber(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_atfa_dev = MagicMock() |
| atft_manager.atfa_dev = mock_atfa_dev |
| test_atfa_device_manager = atftman.AtfaDeviceManager(atft_manager) |
| atft_manager.product_info = MagicMock() |
| atft_manager.product_info.product_id = self.TEST_ID |
| mock_atfa_dev.Oem.return_value = 'TEST\n(bootloader) abcd\nTEST' |
| with self.assertRaises(FastbootFailure): |
| test_atfa_device_manager.CheckStatus() |
| |
| # Test AtftManager.CheckProvisionStatus |
| def MockGetVar(self, variable): |
| return self.status_map.get(variable) |
| |
| def testCheckProvisionStatus(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| # All initial state |
| self.status_map = {} |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked: 0\n' |
| '(bootloader) bootloader-min-versions: -1,0,3\n' |
| '(bootloader) avb-perm-attr-set: 0\n' |
| '(bootloader) avb-locked: 0\n' |
| '(bootloader) avb-unlock-disabled: 0\n' |
| '(bootloader) avb-min-versions: 0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = '' |
| mock_device = MagicMock() |
| mock_device.GetVar.side_effect = self.MockGetVar |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(ProvisionStatus.IDLE, mock_device.provision_status) |
| |
| # Attestation key provisioned |
| self.status_map['at-attest-uuid'] = self.TEST_UUID |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(ProvisionStatus.PROVISION_SUCCESS, |
| mock_device.provision_status) |
| |
| # AVB locked |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked: 0\n' |
| '(bootloader) bootloader-min-versions: -1,0,3\n' |
| '(bootloader) avb-perm-attr-set: 0\n' |
| '(bootloader) avb-locked: 1\n' |
| '(bootloader) avb-unlock-disabled: 0\n' |
| '(bootloader) avb-min-versions: 0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = '' |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(ProvisionStatus.LOCKAVB_SUCCESS, |
| mock_device.provision_status) |
| |
| # Permanent attributes fused |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked: 0\n' |
| '(bootloader) bootloader-min-versions: -1,0,3\n' |
| '(bootloader) avb-perm-attr-set: 1\n' |
| '(bootloader) avb-locked: 0\n' |
| '(bootloader) avb-unlock-disabled: 0\n' |
| '(bootloader) avb-min-versions: 0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = '' |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(ProvisionStatus.FUSEATTR_SUCCESS, |
| mock_device.provision_status) |
| |
| # Bootloader locked |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked: 1\n' |
| '(bootloader) bootloader-min-versions: -1,0,3\n' |
| '(bootloader) avb-perm-attr-set: 0\n' |
| '(bootloader) avb-locked: 0\n' |
| '(bootloader) avb-unlock-disabled: 0\n' |
| '(bootloader) avb-min-versions: 0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = '' |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(ProvisionStatus.FUSEVBOOT_SUCCESS, |
| mock_device.provision_status) |
| |
| def testCheckProvisionStatusFormat(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| # All initial state |
| self.status_map = {} |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked=1\n' |
| '(bootloader) bootloader-min-versions=-1,0,3\n' |
| '(bootloader) avb-perm-attr-set=0\n' |
| '(bootloader) avb-locked=0\n' |
| '(bootloader) avb-unlock-disabled=0\n' |
| '(bootloader) avb-min-versions=0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = '' |
| mock_device = MagicMock() |
| mock_device.GetVar.side_effect = self.MockGetVar |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(ProvisionStatus.FUSEVBOOT_SUCCESS, mock_device.provision_status) |
| self.assertEqual(True, mock_device.provision_state.bootloader_locked) |
| |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked:1\n' |
| '(bootloader) bootloader-min-versions:-1,0,3\n' |
| '(bootloader) avb-perm-attr-set:0\n' |
| '(bootloader) avb-locked:0\n' |
| '(bootloader) avb-unlock-disabled:0\n' |
| '(bootloader) avb-min-versions: 0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = '' |
| mock_device = MagicMock() |
| mock_device.GetVar.side_effect = self.MockGetVar |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(ProvisionStatus.FUSEVBOOT_SUCCESS, mock_device.provision_status) |
| self.assertEqual(True, mock_device.provision_state.bootloader_locked) |
| |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked:\t1\n' |
| '(bootloader) bootloader-min-versions:\t-1,0,3\n' |
| '(bootloader) avb-perm-attr-set:\t0\n' |
| '(bootloader) avb-locked:\t0\n' |
| '(bootloader) avb-unlock-disabled:\t0\n' |
| '(bootloader) avb-min-versions:\t0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = '' |
| mock_device = MagicMock() |
| mock_device.GetVar.side_effect = self.MockGetVar |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(ProvisionStatus.FUSEVBOOT_SUCCESS, mock_device.provision_status) |
| self.assertEqual(True, mock_device.provision_state.bootloader_locked) |
| |
| def testCheckProvisionState(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| # All initial state |
| self.status_map = {} |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked: 0\n' |
| '(bootloader) bootloader-min-versions: -1,0,3\n' |
| '(bootloader) avb-perm-attr-set: 0\n' |
| '(bootloader) avb-locked: 0\n' |
| '(bootloader) avb-unlock-disabled: 0\n' |
| '(bootloader) avb-min-versions: 0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = '' |
| mock_device = MagicMock() |
| mock_device.GetVar.side_effect = self.MockGetVar |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(False, mock_device.provision_state.bootloader_locked) |
| self.assertEqual(False, mock_device.provision_state.avb_perm_attr_set) |
| self.assertEqual(False, mock_device.provision_state.avb_locked) |
| self.assertEqual(False, mock_device.provision_state.provisioned) |
| |
| # Attestation key provisioned |
| self.status_map['at-attest-uuid'] = self.TEST_UUID |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(False, mock_device.provision_state.bootloader_locked) |
| self.assertEqual(False, mock_device.provision_state.avb_perm_attr_set) |
| self.assertEqual(False, mock_device.provision_state.avb_locked) |
| self.assertEqual(True, mock_device.provision_state.provisioned) |
| |
| # AVB locked and attestation key provisioned |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked: 0\n' |
| '(bootloader) bootloader-min-versions: -1,0,3\n' |
| '(bootloader) avb-perm-attr-set: 0\n' |
| '(bootloader) avb-locked: 1\n' |
| '(bootloader) avb-unlock-disabled: 0\n' |
| '(bootloader) avb-min-versions: 0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = self.TEST_UUID |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(False, mock_device.provision_state.bootloader_locked) |
| self.assertEqual(False, mock_device.provision_state.avb_perm_attr_set) |
| self.assertEqual(True, mock_device.provision_state.avb_locked) |
| self.assertEqual(True, mock_device.provision_state.provisioned) |
| self.assertEqual(ProvisionStatus.PROVISION_SUCCESS, |
| mock_device.provision_status) |
| |
| # Permanent attributes fused |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked: 0\n' |
| '(bootloader) bootloader-min-versions: -1,0,3\n' |
| '(bootloader) avb-perm-attr-set: 1\n' |
| '(bootloader) avb-locked: 0\n' |
| '(bootloader) avb-unlock-disabled: 0\n' |
| '(bootloader) avb-min-versions: 0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = '' |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(False, mock_device.provision_state.bootloader_locked) |
| self.assertEqual(True, mock_device.provision_state.avb_perm_attr_set) |
| self.assertEqual(False, mock_device.provision_state.avb_locked) |
| self.assertEqual(False, mock_device.provision_state.provisioned) |
| |
| # All status set |
| self.status_map['at-vboot-state'] = ( |
| '(bootloader) bootloader-locked: 1\n' |
| '(bootloader) bootloader-min-versions: -1,0,3\n' |
| '(bootloader) avb-perm-attr-set: 1\n' |
| '(bootloader) avb-locked: 1\n' |
| '(bootloader) avb-unlock-disabled: 0\n' |
| '(bootloader) avb-min-versions: 0:1,1:1,2:1,4097 :2,4098:2\n') |
| self.status_map['at-attest-uuid'] = self.TEST_UUID |
| atft_manager.CheckProvisionStatus(mock_device) |
| self.assertEqual(True, mock_device.provision_state.bootloader_locked) |
| self.assertEqual(True, mock_device.provision_state.avb_perm_attr_set) |
| self.assertEqual(True, mock_device.provision_state.avb_locked) |
| self.assertEqual(True, mock_device.provision_state.provisioned) |
| self.assertEqual(ProvisionStatus.PROVISION_SUCCESS, |
| mock_device.provision_status) |
| |
| # Test AtftManager.Provision |
| def MockSetProvisionSuccess(self, target): |
| target.provision_status = ProvisionStatus.PROVISION_SUCCESS |
| target.provision_state = ProvisionState() |
| target.provision_state.provisioned = True |
| |
| def MockSetProvisionFail(self, target): |
| target.provision_status = ProvisionStatus.PROVISION_FAILED |
| target.provision_state = ProvisionState() |
| target.provision_state.provisioned = False |
| |
| def testProvision(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_atfa = MagicMock() |
| mock_target = MagicMock() |
| atft_manager._atfa_dev_manager = MagicMock() |
| atft_manager.atfa_dev = mock_atfa |
| atft_manager._GetAlgorithmList = MagicMock() |
| atft_manager._GetAlgorithmList.return_value = [ |
| EncryptionAlgorithm.ALGORITHM_CURVE25519 |
| ] |
| atft_manager.TransferContent = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetProvisionSuccess |
| |
| atft_manager.Provision(mock_target) |
| |
| # Make sure atfa.SetTime is called. |
| atft_manager._atfa_dev_manager.SetTime.assert_called_once() |
| # Transfer content should be ATFA->target, target->ATFA, ATFA->target |
| transfer_content_calls = [ |
| call(mock_atfa, mock_target), |
| call(mock_target, mock_atfa), |
| call(mock_atfa, mock_target) |
| ] |
| atft_manager.TransferContent.assert_has_calls(transfer_content_calls) |
| atfa_oem_calls = [ |
| call('atfa-start-provisioning ' + |
| str(EncryptionAlgorithm.ALGORITHM_CURVE25519)), |
| call('atfa-finish-provisioning') |
| ] |
| target_oem_calls = [ |
| call('at-get-ca-request'), |
| call('at-set-ca-response') |
| ] |
| mock_atfa.Oem.assert_has_calls(atfa_oem_calls) |
| mock_target.Oem.assert_has_calls(target_oem_calls) |
| |
| def testProvisionFailed(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_atfa = MagicMock() |
| mock_target = MagicMock() |
| atft_manager._atfa_dev_manager = MagicMock() |
| atft_manager.atfa_dev = mock_atfa |
| atft_manager._GetAlgorithmList = MagicMock() |
| atft_manager._GetAlgorithmList.return_value = [ |
| EncryptionAlgorithm.ALGORITHM_CURVE25519 |
| ] |
| atft_manager.TransferContent = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetProvisionFail |
| with self.assertRaises(FastbootFailure): |
| atft_manager.Provision(mock_target) |
| |
| # Test AtftManager.FuseVbootKey |
| def MockSetFuseVbootSuccess(self, target): |
| target.provision_status = ProvisionStatus.FUSEVBOOT_SUCCESS |
| target.provision_state = ProvisionState() |
| target.provision_state.bootloader_locked = True |
| |
| def MockSetFuseVbootFail(self, target): |
| target.provision_status = ProvisionStatus.FUSEVBOOT_FAILED |
| target.provision_state = ProvisionState() |
| target.provision_state.bootloader_locked = False |
| |
| @patch('os.remove') |
| @patch('tempfile.NamedTemporaryFile') |
| def testFuseVbootKey(self, mock_create_temp_file, mock_remove): |
| mock_file = MagicMock() |
| mock_create_temp_file.return_value = mock_file |
| mock_file.name = self.TEST_FILE_NAME |
| |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager.product_info = ProductInfo( |
| self.TEST_ID, self.TEST_NAME, self.TEST_ATTRIBUTE_ARRAY, |
| self.TEST_VBOOT_KEY_ARRAY) |
| mock_target = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetFuseVbootSuccess |
| |
| atft_manager.FuseVbootKey(mock_target) |
| |
| mock_file.write.assert_called_once_with(self.TEST_VBOOT_KEY_ARRAY) |
| mock_target.Download.assert_called_once_with(self.TEST_FILE_NAME) |
| mock_remove.assert_called_once_with(self.TEST_FILE_NAME) |
| mock_target.Oem.assert_called_once_with('fuse at-bootloader-vboot-key') |
| |
| @patch('os.remove') |
| @patch('tempfile.NamedTemporaryFile') |
| def testFuseVbootKeyFailed(self, mock_create_temp_file, _): |
| mock_file = MagicMock() |
| mock_create_temp_file.return_value = mock_file |
| mock_file.name = self.TEST_FILE_NAME |
| |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager.product_info = ProductInfo( |
| self.TEST_ID, self.TEST_NAME, self.TEST_ATTRIBUTE_ARRAY, |
| self.TEST_VBOOT_KEY_ARRAY) |
| mock_target = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetFuseVbootFail |
| with self.assertRaises(FastbootFailure): |
| atft_manager.FuseVbootKey(mock_target) |
| |
| def testFuseVbootKeyNoProduct(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_target = MagicMock() |
| atft_manager.product_info = None |
| with self.assertRaises(ProductNotSpecifiedException): |
| atft_manager.FuseVbootKey(mock_target) |
| |
| @patch('os.remove') |
| @patch('tempfile.NamedTemporaryFile') |
| def testFuseVbootKeyFastbootFailure(self, mock_create_temp_file, _): |
| mock_file = MagicMock() |
| mock_create_temp_file.return_value = mock_file |
| mock_file.name = self.TEST_FILE_NAME |
| |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager.product_info = ProductInfo( |
| self.TEST_ID, self.TEST_NAME, self.TEST_ATTRIBUTE_ARRAY, |
| self.TEST_VBOOT_KEY_ARRAY) |
| mock_target = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| mock_target.Oem.side_effect = FastbootFailure('') |
| |
| with self.assertRaises(FastbootFailure): |
| atft_manager.FuseVbootKey(mock_target) |
| self.assertEqual( |
| ProvisionStatus.FUSEVBOOT_FAILED, mock_target.provision_status) |
| |
| # Test AtftManager.FusePermAttr |
| def MockSetFuseAttrSuccess(self, target): |
| target.provision_status = ProvisionStatus.FUSEATTR_SUCCESS |
| target.provision_state = ProvisionState() |
| target.provision_state.avb_perm_attr_set = True |
| |
| def MockSetFuseAttrFail(self, target): |
| target.provision_status = ProvisionStatus.FUSEATTR_FAILED |
| target.provision_state = ProvisionState() |
| target.provision_state.avb_perm_attr_set = False |
| |
| @patch('os.remove') |
| @patch('tempfile.NamedTemporaryFile') |
| def testFusePermAttr(self, mock_create_temp_file, mock_remove): |
| mock_file = MagicMock() |
| mock_create_temp_file.return_value = mock_file |
| mock_file.name = self.TEST_FILE_NAME |
| |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager.product_info = ProductInfo( |
| self.TEST_ID, self.TEST_NAME, self.TEST_ATTRIBUTE_ARRAY, |
| self.TEST_VBOOT_KEY_ARRAY) |
| mock_target = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetFuseAttrSuccess |
| |
| atft_manager.FusePermAttr(mock_target) |
| |
| mock_file.write.assert_called_once_with(self.TEST_ATTRIBUTE_ARRAY) |
| mock_target.Download.assert_called_once_with(self.TEST_FILE_NAME) |
| mock_remove.assert_called_once_with(self.TEST_FILE_NAME) |
| mock_target.Oem.assert_called_once_with('fuse at-perm-attr') |
| |
| @patch('os.remove') |
| @patch('tempfile.NamedTemporaryFile') |
| def testFusePermAttrFail(self, mock_create_temp_file, _): |
| mock_file = MagicMock() |
| mock_create_temp_file.return_value = mock_file |
| mock_file.name = self.TEST_FILE_NAME |
| |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager.product_info = ProductInfo( |
| self.TEST_ID, self.TEST_NAME, self.TEST_ATTRIBUTE_ARRAY, |
| self.TEST_VBOOT_KEY_ARRAY) |
| mock_target = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetFuseAttrFail |
| with self.assertRaises(FastbootFailure): |
| atft_manager.FusePermAttr(mock_target) |
| |
| def testFusePermAttrNoProduct(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_target = MagicMock() |
| atft_manager.product_info = None |
| with self.assertRaises(ProductNotSpecifiedException): |
| atft_manager.FusePermAttr(mock_target) |
| |
| @patch('os.remove') |
| @patch('tempfile.NamedTemporaryFile') |
| def testFusePermAttrFastbootFailure(self, mock_create_temp_file, _): |
| mock_file = MagicMock() |
| mock_create_temp_file.return_value = mock_file |
| mock_file.name = self.TEST_FILE_NAME |
| |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager.product_info = ProductInfo( |
| self.TEST_ID, self.TEST_NAME, self.TEST_ATTRIBUTE_ARRAY, |
| self.TEST_VBOOT_KEY_ARRAY) |
| mock_target = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| mock_target.Oem.side_effect = FastbootFailure('') |
| |
| with self.assertRaises(FastbootFailure): |
| atft_manager.FusePermAttr(mock_target) |
| self.assertEqual( |
| ProvisionStatus.FUSEATTR_FAILED, mock_target.provision_status) |
| |
| # Test AtftManager.LockAvb |
| def MockSetLockAvbSuccess(self, target): |
| target.provision_status = ProvisionStatus.LOCKAVB_SUCCESS |
| target.provision_state = ProvisionState() |
| target.provision_state.avb_locked = True |
| |
| def MockSetLockAvbFail(self, target): |
| target.provision_status = ProvisionStatus.LOCKAVB_FAILED |
| target.provision_state = ProvisionState() |
| target.provision_state.avb_locked = False |
| |
| def testLockAvb(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_target = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetLockAvbSuccess |
| atft_manager.LockAvb(mock_target) |
| mock_target.Oem.assert_called_once_with('at-lock-vboot') |
| self.assertEqual( |
| ProvisionStatus.LOCKAVB_SUCCESS, mock_target.provision_status) |
| |
| def testLockAvbFail(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_target = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetLockAvbFail |
| with self.assertRaises(FastbootFailure): |
| atft_manager.LockAvb(mock_target) |
| |
| def testLockAvbFastbootFailure(self): |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_target = MagicMock() |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetLockAvbSuccess |
| mock_target.Oem.side_effect = FastbootFailure('') |
| with self.assertRaises(FastbootFailure): |
| atft_manager.LockAvb(mock_target) |
| self.assertEqual( |
| ProvisionStatus.LOCKAVB_FAILED, mock_target.provision_status) |
| |
| # Test AtftManager.Reboot |
| class MockTimer(object): |
| def __init__(self, interval, callback): |
| self.interval = interval |
| self.callback = callback |
| |
| def start(self): |
| pass |
| |
| def refresh(self): |
| if self.callback: |
| self.callback() |
| |
| def cancel(self): |
| self.callback = None |
| |
| def mock_create_timer(self, interval, callback): |
| self.mock_timer_instance = self.MockTimer(interval, callback) |
| return self.mock_timer_instance |
| |
| @patch('threading.Timer') |
| def testRebootSuccess(self, mock_timer): |
| self.mock_timer_instance = None |
| atft_manager = atftman.AtftManager( |
| self.FastbootDeviceTemplate, self.mock_serial_mapper, self.configs) |
| timeout = 1 |
| atft_manager.stable_serials = [self.TEST_SERIAL] |
| mock_fastboot = MagicMock() |
| test_device = atftman.DeviceInfo( |
| mock_fastboot, self.TEST_SERIAL, self.TEST_LOCATION) |
| |
| atft_manager.target_devs.append(test_device) |
| mock_success = MagicMock() |
| mock_fail = MagicMock() |
| mock_timer.side_effect = self.mock_create_timer |
| |
| atft_manager.Reboot(test_device, timeout, mock_success, mock_fail) |
| |
| # During the reboot, the status should be REBOOT_ING. |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| self.assertEqual( |
| ProvisionStatus.REBOOT_ING, |
| atft_manager.target_devs[0].provision_status) |
| self.assertEqual( |
| self.TEST_SERIAL, atft_manager.target_devs[0].serial_number) |
| |
| # After the device reappear, the status should be REBOOT_SUCCESS. |
| atft_manager.stable_serials = [self.TEST_SERIAL] |
| atft_manager._HandleRebootCallbacks() |
| # mock timeout event. |
| self.mock_timer_instance.refresh() |
| |
| self.assertEqual(1, len(atft_manager.target_devs)) |
| self.assertEqual( |
| ProvisionStatus.REBOOT_SUCCESS, |
| atft_manager.target_devs[0].provision_status) |
| mock_fastboot.Reboot.assert_called_once() |
| |
| # Success should be called, fail should not. |
| mock_success.assert_called_once() |
| mock_fail.assert_not_called() |
| |
| @patch('threading.Timer') |
| def testRebootTimeout(self, mock_timer): |
| self.mock_timer_instance = None |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| timeout = 1 |
| atft_manager.stable_serials.append(self.TEST_SERIAL) |
| mock_fastboot = MagicMock() |
| test_device = atftman.DeviceInfo( |
| mock_fastboot, self.TEST_SERIAL, self.TEST_LOCATION) |
| atft_manager.target_devs.append(test_device) |
| mock_success = MagicMock() |
| mock_fail = MagicMock() |
| # Status would be checked after reboot. We assume it's in FUSEVBOOT_SUCCESS |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetFuseVbootSuccess |
| mock_timer.side_effect = self.mock_create_timer |
| |
| atft_manager.Reboot(test_device, timeout, mock_success, mock_fail) |
| atft_manager.stable_serials = [] |
| |
| atft_manager._HandleRebootCallbacks() |
| |
| # mock timeout event. |
| self.mock_timer_instance.refresh() |
| |
| self.assertEqual(0, len(atft_manager.target_devs)) |
| mock_success.assert_not_called() |
| mock_fail.assert_called_once() |
| |
| @patch('threading.Timer') |
| def testRebootTimeoutBeforeRefresh(self, mock_timer): |
| self.mock_timer_instance = None |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| timeout = 1 |
| atft_manager.stable_serials.append(self.TEST_SERIAL) |
| mock_fastboot = MagicMock() |
| test_device = atftman.DeviceInfo( |
| mock_fastboot, self.TEST_SERIAL, self.TEST_LOCATION) |
| atft_manager.target_devs.append(test_device) |
| mock_success = MagicMock() |
| mock_fail = MagicMock() |
| # Status would be checked after reboot. We assume it's in FUSEVBOOT_SUCCESS |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetFuseVbootSuccess |
| mock_timer.side_effect = self.mock_create_timer |
| |
| atft_manager.Reboot(test_device, timeout, mock_success, mock_fail) |
| atft_manager.stable_serials = [] |
| # mock timeout event. |
| self.mock_timer_instance.refresh() |
| # mock refresh event. |
| atft_manager._HandleRebootCallbacks() |
| |
| self.assertEqual(0, len(atft_manager.target_devs)) |
| mock_success.assert_not_called() |
| mock_fail.assert_called_once() |
| |
| @patch('threading.Timer') |
| def testRebootFailure(self, mock_timer): |
| self.mock_timer_instance = None |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| mock_target = MagicMock() |
| mock_target.serial_number = self.TEST_SERIAL |
| timeout = 1 |
| atft_manager.stable_serials.append(self.TEST_SERIAL) |
| test_device = atftman.DeviceInfo(None, self.TEST_SERIAL, self.TEST_LOCATION) |
| atft_manager.target_devs.append(test_device) |
| mock_success = MagicMock() |
| mock_fail = MagicMock() |
| # Status would be checked after reboot. We assume it's in FUSEVBOOT_SUCCESS |
| atft_manager.CheckProvisionStatus = MagicMock() |
| atft_manager.CheckProvisionStatus.side_effect = self.MockSetFuseVbootSuccess |
| mock_timer.side_effect = self.mock_create_timer |
| mock_target.Reboot.side_effect = FastbootFailure('') |
| |
| with self.assertRaises(FastbootFailure): |
| atft_manager.Reboot(mock_target, timeout, mock_success, mock_fail) |
| |
| # There should be no timeout timer. |
| self.assertEqual(None, self.mock_timer_instance) |
| # mock refresh event. |
| atft_manager._HandleRebootCallbacks() |
| mock_success.assert_not_called() |
| mock_fail.assert_not_called() |
| |
| # Test AtftManager.ProcessProductAttributesFile |
| def testProcessProductAttributesFile(self): |
| test_content = ( |
| '{' |
| ' "productName": "%s",' |
| ' "productConsoleId": "%s",' |
| ' "productPermanentAttribute": "%s",' |
| ' "bootloaderPublicKey": "%s",' |
| ' "creationTime": ""' |
| '}') % (self.TEST_NAME, self.TEST_ID, self.TEST_ATTRIBUTE_STRING, |
| self.TEST_VBOOT_KEY_STRING) |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| atft_manager.ProcessProductAttributesFile(test_content) |
| self.assertEqual(self.TEST_NAME, atft_manager.product_info.product_name) |
| self.assertEqual(self.TEST_ID, atft_manager.product_info.product_id) |
| self.assertEqual(self.TEST_ATTRIBUTE_ARRAY, |
| atft_manager.product_info.product_attributes) |
| self.assertEqual(self.TEST_VBOOT_KEY_ARRAY, |
| atft_manager.product_info.vboot_key) |
| |
| def testProcessProductAttributesFileWrongJSON(self): |
| test_content = ( |
| '{' |
| ' "productName": "%s",' |
| ' "productConsoleId": "%s",' |
| ' "productPermanentAttribute": "%s",' |
| ' "bootloaderPublicKey": "%s",' |
| ' "creationTime": ""' |
| '') % (self.TEST_NAME, self.TEST_ID, self.TEST_ATTRIBUTE_STRING, |
| self.TEST_VBOOT_KEY_STRING) |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| with self.assertRaises(ProductAttributesFileFormatError): |
| atft_manager.ProcessProductAttributesFile(test_content) |
| |
| def testProcessProductAttributesFileMissingField(self): |
| test_content = ( |
| '{' |
| ' "productConsoleId": "%s",' |
| ' "productPermanentAttribute": "%s",' |
| ' "bootloaderPublicKey": "%s",' |
| ' "creationTime": ""' |
| '}') % (self.TEST_ID, self.TEST_ATTRIBUTE_STRING, |
| self.TEST_VBOOT_KEY_STRING) |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| with self.assertRaises(ProductAttributesFileFormatError): |
| atft_manager.ProcessProductAttributesFile(test_content) |
| |
| def testProcessProductAttributesFileWrongLength(self): |
| test_content = ( |
| '{' |
| ' "productName": "%s",' |
| ' "productConsoleId": "%s",' |
| ' "productPermanentAttribute": "%s",' |
| ' "bootloaderPublicKey": "%s",' |
| ' "creationTime": ""' |
| '}') % (self.TEST_NAME, self.TEST_ID, |
| base64.standard_b64encode(bytearray(1053)), |
| self.TEST_VBOOT_KEY_STRING) |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| with self.assertRaises(ProductAttributesFileFormatError) as e: |
| atft_manager.ProcessProductAttributesFile(test_content) |
| |
| def testProcessProductAttributesFileWrongBase64(self): |
| test_content = ( |
| '{' |
| ' "productName": "%s",' |
| ' "productConsoleId": "%s",' |
| ' "productPermanentAttribute": "%s",' |
| ' "bootloaderPublicKey": "%s",' |
| ' "creationTime": ""' |
| '}') % (self.TEST_NAME, self.TEST_ID, self.TEST_ATTRIBUTE_STRING, |
| '12') |
| atft_manager = atftman.AtftManager(self.FastbootDeviceTemplate, |
| self.mock_serial_mapper, self.configs) |
| with self.assertRaises(ProductAttributesFileFormatError): |
| atft_manager.ProcessProductAttributesFile(test_content) |
| |
| if __name__ == '__main__': |
| unittest.main() |
