Remove partially installer image when image install failed. (sonic-net#3712)

Remove partially installer image when image install failed.

What I did
When install image failed, partially installed image not removed, which may cause disk full issue on small disk device.

How I did it
Handle install image SystemExit exception and uninstall image.

How to verify it
Manually verify.
Pass all test case

Author: liuh-80

sonic_installer/main.py

@@ -1,6 +1,7 @@
 import configparser
 import os
 import re
+import shutil
 import subprocess
 import sys
 import time
@@ -591,7 +592,18 @@ def install(url, force, skip_platform_check=False, skip_migration=False, skip_pa
 
         echo_and_log("Installing image {} and setting it as default...".format(binary_image_version))
         with SWAPAllocator(not skip_setup_swap, swap_mem_size, total_mem_threshold, available_mem_threshold):
-            bootloader.install_image(image_path)
+            try:
+                bootloader.install_image(image_path)
+            except SystemExit as e:
+                # When install image failed with exception, image is partial installed
+                # Fully installed image will update SONiC environment by update_sonic_environment
+                # For partial installed image, only need delete image folder
+                echo_and_log('Image install failed with exception: {}'.format(e))
+                echo_and_log('Delete partial installed image: {}'.format(binary_image_version))
+                new_image_dir = bootloader.get_image_path(binary_image_version)
+                shutil.rmtree(new_image_dir)
+                raise
+
         # Take a backup of current configuration
         if skip_migration:
             echo_and_log("Skipping configuration migration as requested in the command option.")

tests/test_sonic_installer.py

@@ -147,3 +147,52 @@ def test_runtime_exception(mock_popen):
     assert '\nSTDERR:\nFailed' in sre.value.notes, "Invalid STDERR"
 
     assert all(v in str(sre.value) for v in ['test.sh', 'Running', 'Failed']), "Invalid message"
+
+
+@patch("sonic_installer.main.SWAPAllocator")
+@patch("sonic_installer.main.get_bootloader")
+@patch("sonic_installer.main.run_command_or_raise")
+@patch("sonic_installer.main.run_command")
+@patch('shutil.rmtree')
+def test_install_failed(rmtree, run_command, run_command_or_raise, get_bootloader, swap, fs):
+    """ This test covers the "sonic-installer" install image failed handling. """
+
+    sonic_image_filename = "sonic.bin"
+    current_image_version = "image_1"
+    new_image_version = "image_2"
+    new_image_folder = f"/images/{new_image_version}"
+    image_docker_folder = os.path.join(new_image_folder, "docker")
+    mounted_image_folder = f"/tmp/image-{new_image_version}-fs"
+    dockerd_opts = ["--iptables=false", "--bip=1.1.1.1/24"]
+
+    # Setup mock files needed for our test scenario
+    fs.create_file(sonic_image_filename)
+    fs.create_dir(image_docker_folder)
+    fs.create_dir(os.path.join(mounted_image_folder, "usr/lib/docker/docker.sh"))
+    fs.create_file("/var/run/docker.pid", contents="15")
+    fs.create_file("/proc/15/cmdline", contents="\x00".join(["dockerd"] + dockerd_opts))
+
+    # Setup bootloader mock
+    mock_bootloader = Mock()
+    mock_bootloader.get_binary_image_version = Mock(return_value=new_image_version)
+    mock_bootloader.get_installed_images = Mock(return_value=[current_image_version])
+    mock_bootloader.get_image_path = Mock(return_value=new_image_folder)
+    mock_bootloader.verify_image_sign = Mock(return_value=True)
+
+    def mock_install_image(arg):
+        sys.exit(1)
+
+    mock_bootloader.install_image = mock_install_image
+
+    @contextmanager
+    def rootfs_path_mock(path):
+        yield mounted_image_folder
+
+    mock_bootloader.get_rootfs_path = rootfs_path_mock
+
+    get_bootloader.return_value = mock_bootloader
+
+    # Invoke CLI command
+    runner = CliRunner()
+    result = runner.invoke(sonic_installer.commands["install"], [sonic_image_filename, "-y"])
+    print(result.output)