test/py/tests: Move "setup_image" from test_ut.py to fs_helper.py

The generic function in test_ut.py to create a disk image with partition
table can be useful outside of test_ut.py so move it to be available
more clearly.

To make this a bit more easily used library function, make use of
check_call directly rather than calling things though u_boot_utils. In
turn, to more easily handle stdin here, use the shell "printf" utility
to pass sfdisk the specification to create as we do not have an actual
file descriptor to use here.

Signed-off-by: Tom Rini <trini@konsulko.com>

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

test/py/tests/test_ut.py

@@ -27,43 +27,13 @@ def mkdir_cond(dirname):
     if not os.path.exists(dirname):
         os.mkdir(dirname)
 
-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')
-    mkdir_cond(mnt)
-
-    spec = f'type={part_type:x}, size={img_size - 2}M, start=1M, bootable'
-    if second_part:
-        spec += '\ntype=c'
-
-    utils.run_and_log(ubman, f'qemu-img create {fname} 20M')
-    utils.run_and_log(ubman, f'sfdisk {fname}',
-                             stdin=spec.encode('utf-8'))
-    return fname, mnt
-
 def setup_bootmenu_image(ubman):
     """Create a 20MB disk image with a single ext4 partition
 
     This is modelled on Armbian 22.08 Jammy
     """
     mmc_dev = 4
-    fname, mnt = setup_image(ubman, mmc_dev, 0x83)
+    fname, mnt = fs_helper.setup_image(ubman, mmc_dev, 0x83)
 
     script = '''# DO NOT EDIT THIS FILE
 #
@@ -179,7 +149,7 @@ booti ${kernel_addr_r} ${ramdisk_addr_r} ${fdt_addr_r}
 def setup_bootflow_image(ubman):
     """Create a 20MB disk image with a single FAT partition"""
     mmc_dev = 1
-    fname, mnt = setup_image(ubman, mmc_dev, 0xc, second_part=True)
+    fname, mnt = fs_helper.setup_image(ubman, mmc_dev, 0xc, second_part=True)
 
     vmlinux = 'vmlinuz-5.3.7-301.fc31.armv7hl'
     initrd = 'initramfs-5.3.7-301.fc31.armv7hl.img'
@@ -545,8 +515,8 @@ def setup_efi_image(ubman):
     """Create a 20MB disk image with an EFI app on it"""
     devnum = 1
     basename = 'flash'
-    fname, mnt = setup_image(ubman, devnum, 0xc, second_part=True,
+    fname, mnt = fs_helper.setup_image(ubman, devnum, 0xc, second_part=True,
-                             basename=basename)
+                                       basename=basename)
 
     efi_dir = os.path.join(mnt, 'EFI')
     mkdir_cond(efi_dir)