Image-installer: Fix duplication of image prefix

sonic_installer/bootloader/aboot.py

@@ -92,27 +92,28 @@ def _boot_config_set(self, **kwargs):
         self._boot_config_write(config, path=path)
 
     def _swi_image_path(self, image):
-        image_dir = image.replace(IMAGE_PREFIX, IMAGE_DIR_PREFIX)
+        image_dir = image.replace(IMAGE_PREFIX, IMAGE_DIR_PREFIX, 1)
         if is_secureboot():
            return 'flash:%s/sonic.swi' % image_dir
         return 'flash:%s/.sonic-boot.swi' % image_dir
 
     def get_current_image(self):
         with open('/proc/cmdline') as f:
             current = re.search(r"loop=/*(\S+)/", f.read()).group(1)
-        return current.replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX)
+        return current.replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX, 1)
 
     def get_installed_images(self):
         images = []
         for filename in os.listdir(HOST_PATH):
             if filename.startswith(IMAGE_DIR_PREFIX):
-                images.append(filename.replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX))
+                images.append(filename.replace(IMAGE_DIR_PREFIX,
+                                               IMAGE_PREFIX, 1))
         return images
 
     def get_next_image(self):
         config = self._boot_config_read()
         match = re.search(r"flash:/*(\S+)/", config['SWI'])
-        return match.group(1).replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX)
+        return match.group(1).replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX, 1)
 
     def set_default_image(self, image):
         image_path = self._swi_image_path(image)

sonic_installer/bootloader/bootloader.py

@@ -76,7 +76,7 @@ def detect(cls):
     def get_image_path(cls, image):
         """returns the image path"""
         prefix = path.join(HOST_PATH, IMAGE_DIR_PREFIX)
-        return image.replace(IMAGE_PREFIX, prefix)
+        return image.replace(IMAGE_PREFIX, prefix, 1)
 
     @contextmanager
     def get_rootfs_path(self, image_path):

sonic_installer/bootloader/grub.py

@@ -77,7 +77,7 @@ def remove_image(self, image):
         config.close()
         click.echo('Done')
 
-        image_dir = image.replace(IMAGE_PREFIX, IMAGE_DIR_PREFIX)
+        image_dir = image.replace(IMAGE_PREFIX, IMAGE_DIR_PREFIX, 1)
         click.echo('Removing image root filesystem...')
         subprocess.call(['rm','-rf', HOST_PATH + '/' + image_dir])
         click.echo('Done')

sonic_installer/bootloader/onie.py

@@ -21,7 +21,10 @@ def get_current_image(self):
         cmdline = open('/proc/cmdline', 'r')
         current = re.search(r"loop=(\S+)/fs.squashfs", cmdline.read()).group(1)
         cmdline.close()
-        return current.replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX)
+        # Replace only the first occurrence, since we are using branch name as
+        # tagging for version, IMAGE_PREFIX in the name of the branch may be
+        # replaced as well.
+        return current.replace(IMAGE_DIR_PREFIX, IMAGE_PREFIX, 1)
 
     def get_binary_image_version(self, image_path):
         """returns the version of the image"""

sonic_installer/bootloader/uboot.py

@@ -73,7 +73,7 @@ def remove_image(self, image):
         elif image in images[1]:
             run_command('/usr/bin/fw_setenv boot_next "run sonic_image_1"')
             run_command('/usr/bin/fw_setenv sonic_version_2 "NONE"')
-        image_dir = image.replace(IMAGE_PREFIX, IMAGE_DIR_PREFIX)
+        image_dir = image.replace(IMAGE_PREFIX, IMAGE_DIR_PREFIX, 1)
         click.echo('Removing image root filesystem...')
         subprocess.call(['rm','-rf', HOST_PATH + '/' + image_dir])
         click.echo('Done')

sonic_installer/main.py

@@ -272,7 +272,7 @@ def update_sonic_environment(bootloader, binary_image_version):
     SONIC_ENV_TEMPLATE_FILE = os.path.join("usr", "share", "sonic", "templates", "sonic-environment.j2")
     SONIC_VERSION_YML_FILE = os.path.join("etc", "sonic", "sonic_version.yml")
 
-    sonic_version = re.sub(IMAGE_PREFIX, '', binary_image_version)
+    sonic_version = re.sub(IMAGE_PREFIX, '', binary_image_version, 1)
     new_image_dir = bootloader.get_image_path(binary_image_version)
     new_image_mount = os.path.join('/', "tmp", "image-{0}-fs".format(sonic_version))
     env_dir = os.path.join(new_image_dir, "sonic-config")
@@ -318,7 +318,7 @@ def migrate_sonic_packages(bootloader, binary_image_version):
     tmp_dir = "tmp"
     packages_file = "packages.json"
     packages_path = os.path.join(PACKAGE_MANAGER_DIR, packages_file)
-    sonic_version = re.sub(IMAGE_PREFIX, '', binary_image_version)
+    sonic_version = re.sub(IMAGE_PREFIX, '', binary_image_version, 1)
     new_image_dir = bootloader.get_image_path(binary_image_version)
     new_image_upper_dir = os.path.join(new_image_dir, UPPERDIR_NAME)
     new_image_work_dir = os.path.join(new_image_dir, WORKDIR_NAME)

tests/installer_bootloader_aboot_test.py

@@ -0,0 +1,52 @@
+from unittest.mock import Mock, patch
+
+# Import test module
+import sonic_installer.bootloader.aboot as aboot
+
+# Constants
+image_dir = f'{aboot.IMAGE_DIR_PREFIX}expeliarmus-{aboot.IMAGE_DIR_PREFIX}abcde'
+exp_image = f'{aboot.IMAGE_PREFIX}expeliarmus-{aboot.IMAGE_DIR_PREFIX}abcde'
+image_dirs = [image_dir]
+
+
+@patch('sonic_installer.bootloader.aboot.is_secureboot',
+       Mock(return_value=False))
+def test_swi_image_path():
+    # Constants
+    image_id = f'{aboot.IMAGE_PREFIX}expeliarmus-{aboot.IMAGE_PREFIX}abcde'
+    exp_image_path = f'flash:{aboot.IMAGE_DIR_PREFIX}expeliarmus-'\
+                     f'{aboot.IMAGE_PREFIX}abcde/.sonic-boot.swi'
+
+    bootloader = aboot.AbootBootloader()
+
+    # Verify converted swi image path
+    image_path = bootloader._swi_image_path(image_id)
+    assert image_path == exp_image_path
+
+
+@patch("sonic_installer.bootloader.aboot.re.search")
+def test_get_current_image(re_search_patch):
+    bootloader = aboot.AbootBootloader()
+
+    # Test convertion image dir to image name
+    re_search_patch().group = Mock(return_value=image_dir)
+    assert bootloader.get_current_image() == exp_image
+
+
+@patch('sonic_installer.bootloader.aboot.os.listdir',
+       Mock(return_value=image_dirs))
+def test_get_installed_images():
+    bootloader = aboot.AbootBootloader()
+
+    # Test convertion image dir to image name
+    assert bootloader.get_installed_images() == [exp_image]
+
+
+@patch("sonic_installer.bootloader.aboot.re.search")
+def test_get_next_image(re_search_patch):
+    bootloader = aboot.AbootBootloader()
+    bootloader._boot_config_read = Mock(return_value={'SWI': None})
+
+    # Test convertion image dir to image name
+    re_search_patch().group = Mock(return_value=image_dir)
+    assert bootloader.get_next_image() == exp_image

tests/installer_bootloader_bootloader_test.py

@@ -0,0 +1,16 @@
+import os
+
+# Import test module
+import sonic_installer.bootloader.bootloader as bl
+
+
+def test_get_image_path():
+    # Constants
+    image = f'{bl.IMAGE_PREFIX}expeliarmus-{bl.IMAGE_PREFIX}abcde'
+    path_prefix = os.path.join(bl.HOST_PATH, bl.IMAGE_DIR_PREFIX)
+    exp_image_path = f'{path_prefix}expeliarmus-{bl.IMAGE_PREFIX}abcde'
+
+    bootloader = bl.Bootloader()
+
+    # Test replacement image id with image path
+    assert bootloader.get_image_path(image) == exp_image_path

tests/installer_bootloader_grub_test.py

@@ -0,0 +1,26 @@
+import os
+from unittest.mock import Mock, patch
+
+# Import test module
+import sonic_installer.bootloader.grub as grub
+
+
+@patch("sonic_installer.bootloader.grub.subprocess.call", Mock())
+@patch("sonic_installer.bootloader.grub.open")
+@patch("sonic_installer.bootloader.grub.run_command")
+@patch("sonic_installer.bootloader.grub.re.search")
+def test_remove_image(open_patch, run_command_patch, re_search_patch):
+    # Constants
+    image_path_prefix = os.path.join(grub.HOST_PATH, grub.IMAGE_DIR_PREFIX)
+    exp_image_path = f'{image_path_prefix}expeliarmus-{grub.IMAGE_PREFIX}abcde'
+    image = f'{grub.IMAGE_PREFIX}expeliarmus-{grub.IMAGE_PREFIX}abcde'
+
+    bootloader = grub.GrubBootloader()
+
+    # Verify rm command was executed with image path
+    bootloader.remove_image(image)
+    args_list = grub.subprocess.call.call_args_list
+    assert len(args_list) > 0
+
+    args, _ = args_list[0]
+    assert exp_image_path in args[0]

tests/installer_bootloader_onie_test.py

@@ -0,0 +1,17 @@
+from unittest.mock import Mock, patch
+
+# Import test module
+import sonic_installer.bootloader.onie as onie
+
+
+@patch("sonic_installer.bootloader.onie.re.search")
+def test_get_current_image(re_search):
+    # Constants
+    image = f'{onie.IMAGE_DIR_PREFIX}expeliarmus-{onie.IMAGE_DIR_PREFIX}abcde'
+    exp_image = f'{onie.IMAGE_PREFIX}expeliarmus-{onie.IMAGE_DIR_PREFIX}abcde'
+
+    bootloader = onie.OnieInstallerBootloader()
+
+    # Test image dir conversion
+    onie.re.search().group = Mock(return_value=image)
+    assert bootloader.get_current_image() == exp_image

tests/installer_bootloader_uboot_test.py

@@ -0,0 +1,29 @@
+import os
+from unittest.mock import Mock, patch
+
+# Import test module
+import sonic_installer.bootloader.uboot as uboot
+
+
+@patch("sonic_installer.bootloader.uboot.subprocess.call", Mock())
+@patch("sonic_installer.bootloader.uboot.run_command")
+def test_remove_image(run_command_patch):
+    # Constants
+    image_path_prefix = os.path.join(uboot.HOST_PATH, uboot.IMAGE_DIR_PREFIX)
+    exp_image_path = f'{image_path_prefix}expeliarmus-{uboot.IMAGE_PREFIX}abcde'
+
+    intstalled_images = [
+        f'{uboot.IMAGE_PREFIX}expeliarmus-{uboot.IMAGE_PREFIX}abcde',
+        f'{uboot.IMAGE_PREFIX}expeliarmus-abcde',
+    ]
+
+    bootloader = uboot.UbootBootloader()
+    bootloader.get_installed_images = Mock(return_value=intstalled_images)
+
+    # Verify rm command was executed with image path
+    bootloader.remove_image(intstalled_images[0])
+    args_list = uboot.subprocess.call.call_args_list
+    assert len(args_list) > 0
+
+    args, _ = args_list[0]
+    assert exp_image_path in args[0]