1 files changed, 34 insertions, 0 deletions
diff --git a/test/py/tests/fs_helper.py b/test/py/tests/fs_helper.py
index ccfc0201a49..f3e81b6bc61 100644
--- a/test/py/tests/fs_helper.py
+++ b/test/py/tests/fs_helper.py
@@ -69,6 +69,40 @@ def mk_fs(config, fs_type, size, prefix, src_dir=None, size_gran = 0x100000):
         call(f'rm -f {fs_img}', shell=True)
         raise
 
+def setup_image(ubman, devnum, part_type, img_size=20, second_part=False,
+                basename='mmc'):
+    """Create a disk image with a single partition
+
+    Args:
+        ubman (ConsoleBase): Console to use
+        devnum (int): Device number to use, e.g. 1
+        part_type (int): Partition type, e.g. 0xc for FAT32
+        img_size (int): Image size in MiB
+        second_part (bool): True to contain a small second partition
+        basename (str): Base name to use in the filename, e.g. 'mmc'
+
+    Returns:
+        tuple:
+            str: Filename of MMC image
+            str: Directory name of scratch directory
+    """
+    fname = os.path.join(ubman.config.source_dir, f'{basename}{devnum}.img')
+    mnt = os.path.join(ubman.config.persistent_data_dir, 'scratch')
+
+    spec = f'type={part_type:x}, size={img_size - 2}M, start=1M, bootable'
+    if second_part:
+        spec += '\ntype=c'
+
+    try:
+        check_call(f'mkdir -p {mnt}', shell=True)
+        check_call(f'qemu-img create {fname} 20M', shell=True)
+        check_call(f'printf "{spec}" | sfdisk {fname}', shell=True)
+    except CalledProcessError:
+        call(f'rm -f {fname}', shell=True)
+        raise
+
+    return fname, mnt
+
 # Just for trying out
 if __name__ == "__main__":
     import collections