1 files changed, 50 insertions, 3 deletions

diff --git a/test/py/tests/fs_helper.py b/test/py/tests/fs_helper.py
index 380f4c4dca3..b8a22b22806 100644
--- a/test/py/tests/fs_helper.py
+++ b/test/py/tests/fs_helper.py
@@ -9,7 +9,7 @@ import re
 import os
 from subprocess import call, check_call, check_output, CalledProcessError
 
-def mk_fs(config, fs_type, size, prefix, size_gran = 0x100000):
+def mk_fs(config, fs_type, size, prefix, src_dir=None, size_gran = 0x100000):
     """Create a file system volume
 
     Args:
@@ -17,6 +17,7 @@ def mk_fs(config, fs_type, size, prefix, size_gran = 0x100000):
         fs_type (str): File system type, e.g. 'ext4'
         size (int): Size of file system in bytes
         prefix (str): Prefix string of volume's file name
+        src_dir (str): Root directory to use, or None for none
         size_gran (int): Size granularity of file system image in bytes
 
     Raises:
@@ -34,11 +35,19 @@ def mk_fs(config, fs_type, size, prefix, size_gran = 0x100000):
     else:
         mkfs_opt = ''
 
-    if re.match('fat', fs_type):
+    if fs_type == 'exfat':
+        fs_lnxtype = 'exfat'
+    elif re.match('fat', fs_type) or fs_type == 'fs_generic':
         fs_lnxtype = 'vfat'
     else:
         fs_lnxtype = fs_type
 
+    if src_dir:
+        if fs_lnxtype == 'ext4':
+            mkfs_opt = mkfs_opt + ' -d ' + src_dir
+        elif fs_lnxtype != 'vfat' and fs_lnxtype != 'exfat':
+            raise ValueError(f'src_dir not implemented for fs {fs_lnxtype}')
+
     count = (size + size_gran - 1) // size_gran
 
     # Some distributions do not add /sbin to the default PATH, where mkfs lives
@@ -47,7 +56,7 @@ def mk_fs(config, fs_type, size, prefix, size_gran = 0x100000):
 
     try:
         check_call(f'rm -f {fs_img}', shell=True)
-        check_call(f'dd if=/dev/zero of={fs_img} bs={size_gran} count={count}',
+        check_call(f'truncate -s $(( {size_gran} * {count} )) {fs_img}',
                    shell=True)
         check_call(f'mkfs.{fs_lnxtype} {mkfs_opt} {fs_img}', shell=True)
         if fs_type == 'ext4':
@@ -55,11 +64,49 @@ def mk_fs(config, fs_type, size, prefix, size_gran = 0x100000):
                                       shell=True).decode()
             if 'metadata_csum' in sb_content:
                 check_call(f'tune2fs -O ^metadata_csum {fs_img}', shell=True)
+        elif fs_lnxtype == 'vfat' and src_dir:
+            check_call(f'mcopy -i {fs_img} -vsmpQ {src_dir}/* ::/', shell=True)
+        elif fs_lnxtype == 'exfat' and src_dir:
+            check_call(f'fattools cp {src_dir}/* {fs_img}', shell=True)
         return fs_img
     except CalledProcessError:
         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} {img_size}M', 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