image-tools: T6016: wait for umount in cleanup function

2 files changed, 18 insertions, 2 deletions

diff --git a/python/vyos/system/disk.py b/python/vyos/system/disk.py
index 7860d719f..c8908cd5c 100644
--- a/python/vyos/system/disk.py
+++ b/python/vyos/system/disk.py
@@ -16,6 +16,7 @@
 from json import loads as json_loads
 from os import sync
 from dataclasses import dataclass
+from time import sleep
 
 from psutil import disk_partitions
 
@@ -207,13 +208,25 @@ def find_device(mountpoint: str) -> str:
     Returns:
         str: Path to device, Empty if not found
     """
-    mounted_partitions = disk_partitions()
+    mounted_partitions = disk_partitions(all=True)
     for partition in mounted_partitions:
         if partition.mountpoint == mountpoint:
             return partition.mountpoint
     return ''
 
 
+def wait_for_umount(mountpoint: str = '') -> None:
+    """Wait (within reason) for umount to complete
+    """
+    i = 0
+    while find_device(mountpoint):
+        i += 1
+        if i == 5:
+            print(f'Warning: {mountpoint} still mounted')
+            break
+        sleep(1)
+
+
 def disks_size() -> dict[str, int]:
     """Get a dictionary with physical disks and their sizes
 
diff --git a/src/op_mode/image_installer.py b/src/op_mode/image_installer.py
index 501e9b804..d677c2cf8 100755
--- a/src/op_mode/image_installer.py
+++ b/src/op_mode/image_installer.py
@@ -599,6 +599,8 @@ def cleanup(mounts: list[str] = [], remove_items: list[str] = []) -> None:
         print('Unmounting target filesystems')
         for mountpoint in mounts:
             disk.partition_umount(mountpoint)
+        for mountpoint in mounts:
+            disk.wait_for_umount(mountpoint)
     if remove_items:
         print('Removing temporary files')
         for remove_item in remove_items:
@@ -606,7 +608,8 @@ def cleanup(mounts: list[str] = [], remove_items: list[str] = []) -> None:
                 if Path(remove_item).is_file():
                     Path(remove_item).unlink()
                 if Path(remove_item).is_dir():
-                    rmtree(remove_item)
+                    rmtree(remove_item, ignore_errors=True)
+
 
 def cleanup_raid(details: raid.RaidDetails) -> None:
     efiparts = []