sonic-installer: enhance next image detection for Aboot (sonic-net#3433)

The Aboot bootloader relies of the SWI= keyword argument in the
boot-config file to know which image to boot.
This value is also used by sonic-installer to figure to extract the next
image that will be executed.
The current code has an issue as it only expects the next image to match
the installation path of a SONiC image but not anything else.

This means that `SWI=flash:sonic-aboot-broadcom.swi` is not valid and
can therefore be a problem when trying to install a new image via cold
reboot.

Additionally a missing or empty boot-config would generate a python
backtrace instead of gracefully recovering from this state.

Author: Staphylo

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):

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()