diff --git a/sonic_installer/bootloader/aboot.py b/sonic_installer/bootloader/aboot.py index ac327feb4c..d6492171ab 100644 --- a/sonic_installer/bootloader/aboot.py +++ b/sonic_installer/bootloader/aboot.py @@ -71,6 +71,8 @@ class AbootBootloader(Bootloader): def _boot_config_read(self, path=BOOT_CONFIG_PATH): config = collections.OrderedDict() + if not os.path.exists(path): + return config with open(path) as f: for line in f.readlines(): line = line.strip() @@ -112,7 +114,10 @@ def get_installed_images(self): def get_next_image(self): config = self._boot_config_read() - match = re.search(r"flash:/*(\S+)/", config['SWI']) + swi = config.get('SWI', '') + match = re.search(r"flash:/*(\S+)/", swi) + if not match: + return swi.split(':', 1)[-1] return match.group(1).replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX, 1) def set_default_image(self, image): diff --git a/tests/installer_bootloader_aboot_test.py b/tests/installer_bootloader_aboot_test.py index fbe580a638..be09223b5f 100644 --- a/tests/installer_bootloader_aboot_test.py +++ b/tests/installer_bootloader_aboot_test.py @@ -8,6 +8,7 @@ # Constants image_dir = f'{aboot.IMAGE_DIR_PREFIX}expeliarmus-{aboot.IMAGE_DIR_PREFIX}abcde' +image_chainloader = f'{image_dir}/.sonic-boot.swi' exp_image = f'{aboot.IMAGE_PREFIX}expeliarmus-{aboot.IMAGE_DIR_PREFIX}abcde' image_dirs = [image_dir] @@ -45,15 +46,27 @@ def test_get_installed_images(): assert bootloader.get_installed_images() == [exp_image] -@patch("sonic_installer.bootloader.aboot.re.search") -def test_get_next_image(re_search_patch): +def test_get_next_image(): bootloader = aboot.AbootBootloader() - bootloader._boot_config_read = Mock(return_value={'SWI': None}) + + # Test missing boot-config + bootloader._boot_config_read() + + # Test missing SWI value + bootloader._boot_config_read = Mock(return_value={}) + assert bootloader.get_next_image() == '' # Test convertion image dir to image name - re_search_patch().group = Mock(return_value=image_dir) + swi = f'flash:{image_chainloader}' + bootloader._boot_config_read = Mock(return_value={'SWI': swi}) assert bootloader.get_next_image() == exp_image + # Test some other image + next_image = 'EOS.swi' + bootloader._boot_config_read = Mock(return_value={'SWI': f'flash:{next_image}'}) + assert bootloader.get_next_image() == next_image + + def test_install_image(): image_path = 'sonic' env = os.environ.copy()