test: Move common FIT code into a separate fit_util file

To avoid duplicating code, create a new fit_util module which provides
various utility functions for FIT. Move this code out from the existing
test_fit.py and refactor it with addition parameters.

Fix up pylint warnings in the conversion.

This involves no functional change.

Signed-off-by: Simon Glass <sjg@chromium.org>

3 files changed, 110 insertions, 80 deletions

diff --git a/test/boot/bootstd_common.c b/test/boot/bootstd_common.c
index 59d46bef0c5..7a40836507a 100644
--- a/test/boot/bootstd_common.c
+++ b/test/boot/bootstd_common.c
@@ -17,6 +17,7 @@
 #include <u-boot/crc.h>
 #include "bootstd_common.h"
 
+/* tracks whether bootstd_setup_for_tests() has been run yet */
 bool vbe_setup_done;
 
 /* set up MMC for VBE tests */
@@ -27,6 +28,9 @@ int bootstd_setup_for_tests(void)
 	struct blk_desc *desc;
 	int ret;
 
+	if (vbe_setup_done)
+		return 0;
+
 	/* Set up the version string */
 	ret = uclass_get_device(UCLASS_MMC, 1, &mmc);
 	if (ret)
@@ -46,6 +50,8 @@ int bootstd_setup_for_tests(void)
 	if (blk_dwrite(desc, NVDATA_START_BLK, 1, buf) != 1)
 		return log_msg_ret("wr2", -EIO);
 
+	vbe_setup_done = true;
+
 	return 0;
 }
 
@@ -65,17 +71,12 @@ int do_ut_bootstd(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
 	struct unit_test *tests = UNIT_TEST_SUITE_START(bootstd_test);
 	const int n_ents = UNIT_TEST_SUITE_COUNT(bootstd_test);
+	int ret;
 
-	if (!vbe_setup_done) {
-		int ret;
-
-		ret = bootstd_setup_for_tests();
-		if (ret) {
-			printf("Failed to set up for bootstd tests (err=%d)\n",
-			       ret);
-			return CMD_RET_FAILURE;
-		}
-		vbe_setup_done = true;
+	ret = bootstd_setup_for_tests();
+	if (ret) {
+		printf("Failed to set up for bootstd tests (err=%d)\n", ret);
+		return CMD_RET_FAILURE;
 	}
 
 	return cmd_ut_category("bootstd", "bootstd_test_",
diff --git a/test/py/tests/fit_util.py b/test/py/tests/fit_util.py
new file mode 100644
index 00000000000..fcec4c43c3c
--- /dev/null
+++ b/test/py/tests/fit_util.py
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: GPL-2.0+
+# Copyright 2022 Google LLC
+
+"""Common utility functions for FIT tests"""
+
+import os
+
+import u_boot_utils as util
+
+def make_fname(cons, basename):
+    """Make a temporary filename
+
+    Args:
+        cons (ConsoleBase): u_boot_console to use
+        basename (str): Base name of file to create (within temporary directory)
+    Return:
+        Temporary filename
+    """
+
+    return os.path.join(cons.config.build_dir, basename)
+
+def make_its(cons, base_its, params, basename='test.its'):
+    """Make a sample .its file with parameters embedded
+
+    Args:
+        cons (ConsoleBase): u_boot_console to use
+        base_its (str): Template text for the .its file, typically containing
+            %() references
+        params (dict of str): Parameters to embed in the %() strings
+        basename (str): base name to write to (will be placed in the temp dir)
+    Returns:
+        str: Filename of .its file created
+    """
+    its = make_fname(cons, basename)
+    with open(its, 'w', encoding='utf-8') as outf:
+        print(base_its % params, file=outf)
+    return its
+
+def make_fit(cons, mkimage, base_its, params, basename='test.fit'):
+    """Make a sample .fit file ready for loading
+
+    This creates a .its script with the selected parameters and uses mkimage to
+    turn this into a .fit image.
+
+    Args:
+        cons (ConsoleBase): u_boot_console to use
+        mkimage (str): Filename of 'mkimage' utility
+        base_its (str): Template text for the .its file, typically containing
+            %() references
+        params (dict of str): Parameters to embed in the %() strings
+        basename (str): base name to write to (will be placed in the temp dir)
+    Return:
+        Filename of .fit file created
+    """
+    fit = make_fname(cons, basename)
+    its = make_its(cons, base_its, params)
+    util.run_and_log(cons, [mkimage, '-f', its, fit])
+    return fit
+
+def make_kernel(cons, basename, text):
+    """Make a sample kernel with test data
+
+    Args:
+        cons (ConsoleBase): u_boot_console to use
+        basename (str): base name to write to (will be placed in the temp dir)
+        text (str): Contents of the kernel file (will be repeated 100 times)
+    Returns:
+        str: Full path and filename of the kernel it created
+    """
+    fname = make_fname(cons, basename)
+    data = ''
+    for i in range(100):
+        data += f'this {text} {i} is unlikely to boot\n'
+    with open(fname, 'w', encoding='utf-8') as outf:
+        print(data, file=outf)
+    return fname
+
+def make_dtb(cons, base_fdt, basename):
+    """Make a sample .dts file and compile it to a .dtb
+
+    Returns:
+        cons (ConsoleBase): u_boot_console to use
+        Filename of .dtb file created
+    """
+    src = make_fname(cons, f'{basename}.dts')
+    dtb = make_fname(cons, f'{basename}.dtb')
+    with open(src, 'w', encoding='utf-8') as outf:
+        outf.write(base_fdt)
+    util.run_and_log(cons, ['dtc', src, '-O', 'dtb', '-o', dtb])
+    return dtb
diff --git a/test/py/tests/test_fit.py b/test/py/tests/test_fit.py
index 5856960be23..f45848484eb 100755
--- a/test/py/tests/test_fit.py
+++ b/test/py/tests/test_fit.py
@@ -7,6 +7,7 @@ import os
 import pytest
 import struct
 import u_boot_utils as util
+import fit_util
 
 # Define a base ITS which we can adjust using % and a dictionary
 base_its = '''
@@ -126,7 +127,6 @@ def test_fit(u_boot_console):
         Return:
             Temporary filename
         """
-
         return os.path.join(cons.config.build_dir, leaf)
 
     def filesize(fname):
@@ -150,67 +150,6 @@ def test_fit(u_boot_console):
         with open(fname, 'rb') as fd:
             return fd.read()
 
-    def make_dtb():
-        """Make a sample .dts file and compile it to a .dtb
-
-        Returns:
-            Filename of .dtb file created
-        """
-        src = make_fname('u-boot.dts')
-        dtb = make_fname('u-boot.dtb')
-        with open(src, 'w') as fd:
-            fd.write(base_fdt)
-        util.run_and_log(cons, ['dtc', src, '-O', 'dtb', '-o', dtb])
-        return dtb
-
-    def make_its(params):
-        """Make a sample .its file with parameters embedded
-
-        Args:
-            params: Dictionary containing parameters to embed in the %() strings
-        Returns:
-            Filename of .its file created
-        """
-        its = make_fname('test.its')
-        with open(its, 'w') as fd:
-            print(base_its % params, file=fd)
-        return its
-
-    def make_fit(mkimage, params):
-        """Make a sample .fit file ready for loading
-
-        This creates a .its script with the selected parameters and uses mkimage to
-        turn this into a .fit image.
-
-        Args:
-            mkimage: Filename of 'mkimage' utility
-            params: Dictionary containing parameters to embed in the %() strings
-        Return:
-            Filename of .fit file created
-        """
-        fit = make_fname('test.fit')
-        its = make_its(params)
-        util.run_and_log(cons, [mkimage, '-f', its, fit])
-        with open(make_fname('u-boot.dts'), 'w') as fd:
-            fd.write(base_fdt)
-        return fit
-
-    def make_kernel(filename, text):
-        """Make a sample kernel with test data
-
-        Args:
-            filename: the name of the file you want to create
-        Returns:
-            Full path and filename of the kernel it created
-        """
-        fname = make_fname(filename)
-        data = ''
-        for i in range(100):
-            data += 'this %s %d is unlikely to boot\n' % (text, i)
-        with open(fname, 'w') as fd:
-            print(data, file=fd)
-        return fname
-
     def make_ramdisk(filename, text):
         """Make a sample ramdisk with test data
 
@@ -321,10 +260,10 @@ def test_fit(u_boot_console):
           - run code coverage to make sure we are testing all the code
         """
         # Set up invariant files
-        control_dtb = make_dtb()
-        kernel = make_kernel('test-kernel.bin', 'kernel')
+        control_dtb = fit_util.make_dtb(cons, base_fdt, 'u-boot')
+        kernel = fit_util.make_kernel(cons, 'test-kernel.bin', 'kernel')
         ramdisk = make_ramdisk('test-ramdisk.bin', 'ramdisk')
-        loadables1 = make_kernel('test-loadables1.bin', 'lenrek')
+        loadables1 = fit_util.make_kernel(cons, 'test-loadables1.bin', 'lenrek')
         loadables2 = make_ramdisk('test-loadables2.bin', 'ksidmar')
         kernel_out = make_fname('kernel-out.bin')
         fdt = make_fname('u-boot.dtb')
@@ -372,7 +311,7 @@ def test_fit(u_boot_console):
         }
 
         # Make a basic FIT and a script to load it
-        fit = make_fit(mkimage, params)
+        fit = fit_util.make_fit(cons, mkimage, base_its, params)
         params['fit'] = fit
         cmd = base_script % params
 
@@ -403,7 +342,7 @@ def test_fit(u_boot_console):
         # Now a kernel and an FDT
         with cons.log.section('Kernel + FDT load'):
             params['fdt_load'] = 'load = <%#x>;' % params['fdt_addr']
-            fit = make_fit(mkimage, params)
+            fit = fit_util.make_fit(cons, mkimage, base_its, params)
             cons.restart_uboot()
             output = cons.run_command_list(cmd.splitlines())
             check_equal(kernel, kernel_out, 'Kernel not loaded')
@@ -415,7 +354,7 @@ def test_fit(u_boot_console):
         with cons.log.section('Kernel + FDT + Ramdisk load'):
             params['ramdisk_config'] = 'ramdisk = "ramdisk-1";'
             params['ramdisk_load'] = 'load = <%#x>;' % params['ramdisk_addr']
-            fit = make_fit(mkimage, params)
+            fit = fit_util.make_fit(cons, mkimage, base_its, params)
             cons.restart_uboot()
             output = cons.run_command_list(cmd.splitlines())
             check_equal(ramdisk, ramdisk_out, 'Ramdisk not loaded')
@@ -427,7 +366,7 @@ def test_fit(u_boot_console):
                                          params['loadables1_addr'])
             params['loadables2_load'] = ('load = <%#x>;' %
                                          params['loadables2_addr'])
-            fit = make_fit(mkimage, params)
+            fit = fit_util.make_fit(cons, mkimage, base_its, params)
             cons.restart_uboot()
             output = cons.run_command_list(cmd.splitlines())
             check_equal(loadables1, loadables1_out,
@@ -441,7 +380,7 @@ def test_fit(u_boot_console):
             params['kernel'] = make_compressed(kernel)
             params['fdt'] = make_compressed(fdt)
             params['ramdisk'] = make_compressed(ramdisk)
-            fit = make_fit(mkimage, params)
+            fit = fit_util.make_fit(cons, mkimage, base_its, params)
             cons.restart_uboot()
             output = cons.run_command_list(cmd.splitlines())
             check_equal(kernel, kernel_out, 'Kernel not loaded')