3 files changed, 102 insertions, 47 deletions
diff --git a/test/py/tests/test_fs/test_squashfs/sqfs_common.py b/test/py/tests/test_fs/test_squashfs/sqfs_common.py
index 2dc344d1b2..c96f92c1d8 100644
--- a/test/py/tests/test_fs/test_squashfs/sqfs_common.py
+++ b/test/py/tests/test_fs/test_squashfs/sqfs_common.py
@@ -5,6 +5,7 @@
 import os
 import random
 import string
+import subprocess
 
 def sqfs_get_random_letters(size):
     letters = []
@@ -19,24 +20,57 @@ def sqfs_generate_file(path, size):
     file.write(content)
     file.close()
 
-# generate image with three files and a symbolic link
-def sqfs_generate_image(cons):
-    src = os.path.join(cons.config.build_dir, "sqfs_src/")
-    dest = os.path.join(cons.config.build_dir, "sqfs")
-    os.mkdir(src)
-    sqfs_generate_file(src + "frag_only", 100)
-    sqfs_generate_file(src + "blks_frag", 5100)
-    sqfs_generate_file(src + "blks_only", 4096)
-    os.symlink("frag_only", src + "sym")
-    os.system("mksquashfs " + src + " " + dest + " -b 4096 -always-use-fragments")
-
-# removes all files created by sqfs_generate_image()
-def sqfs_clean(cons):
-    src = os.path.join(cons.config.build_dir, "sqfs_src/")
-    dest = os.path.join(cons.config.build_dir, "sqfs")
-    os.remove(src + "frag_only")
-    os.remove(src + "blks_frag")
-    os.remove(src + "blks_only")
-    os.remove(src + "sym")
-    os.rmdir(src)
-    os.remove(dest)
+class Compression:
+    def __init__(self, name, files, sizes, block_size = 4096):
+        self.name = name
+        self.files = files
+        self.sizes = sizes
+        self.mksquashfs_opts = " -b " + str(block_size) + " -comp " + self.name
+
+    def add_opt(self, opt):
+        self.mksquashfs_opts += " " + opt
+
+    def gen_image(self, build_dir):
+        src = os.path.join(build_dir, "sqfs_src/")
+        os.mkdir(src)
+        for (f, s) in zip(self.files, self.sizes):
+            sqfs_generate_file(src + f, s)
+
+        # the symbolic link always targets the first file
+        os.symlink(self.files[0], src + "sym")
+
+        sqfs_img = os.path.join(build_dir, "sqfs-" + self.name)
+        i_o = src + " " + sqfs_img
+        opts = self.mksquashfs_opts
+        try:
+            subprocess.run(["mksquashfs " + i_o + opts], shell = True, check = True)
+        except:
+            print("mksquashfs error. Compression type: " + self.name)
+            raise RuntimeError
+
+    def clean_source(self, build_dir):
+        src = os.path.join(build_dir, "sqfs_src/")
+        for f in self.files:
+            os.remove(src + f)
+        os.remove(src + "sym")
+        os.rmdir(src)
+
+    def cleanup(self, build_dir):
+        self.clean_source(build_dir)
+        sqfs_img = os.path.join(build_dir, "sqfs-" + self.name)
+        os.remove(sqfs_img)
+
+files = ["blks_only", "blks_frag", "frag_only"]
+sizes = [4096, 5100, 100]
+gzip = Compression("gzip", files, sizes)
+zstd = Compression("zstd", files, sizes)
+lzo = Compression("lzo", files, sizes)
+
+# use fragment blocks for files larger than block_size
+gzip.add_opt("-always-use-fragments")
+zstd.add_opt("-always-use-fragments")
+
+# avoid fragments if lzo is used
+lzo.add_opt("-no-fragments")
+
+comp_opts = [gzip, zstd, lzo]
diff --git a/test/py/tests/test_fs/test_squashfs/test_sqfs_load.py b/test/py/tests/test_fs/test_squashfs/test_sqfs_load.py
index eb1baae5c5..9e90062384 100644
--- a/test/py/tests/test_fs/test_squashfs/test_sqfs_load.py
+++ b/test/py/tests/test_fs/test_squashfs/test_sqfs_load.py
@@ -12,23 +12,35 @@ from sqfs_common import *
 @pytest.mark.buildconfigspec('fs_squashfs')
 @pytest.mark.requiredtool('mksquashfs')
 def test_sqfs_load(u_boot_console):
-    cons = u_boot_console
-    sqfs_generate_image(cons)
+    build_dir = u_boot_console.config.build_dir
     command = "sqfsload host 0 $kernel_addr_r "
-    path = os.path.join(cons.config.build_dir, "sqfs")
 
-    try:
+    for opt in comp_opts:
+        # generate and load the squashfs image
+        try:
+            opt.gen_image(build_dir)
+        except RuntimeError:
+            opt.clean_source(build_dir)
+            # skip unsupported compression types
+            continue
+
+        path = os.path.join(build_dir, "sqfs-" + opt.name)
         output = u_boot_console.run_command("host bind 0 " + path)
+
         output = u_boot_console.run_command(command + "xxx")
         assert "File not found." in output
-        output = u_boot_console.run_command(command + "frag_only")
-        assert "100 bytes read in" in output
-        output = u_boot_console.run_command(command + "blks_frag")
-        assert "5100 bytes read in" in output
-        output = u_boot_console.run_command(command + "blks_only")
-        assert "4096 bytes read in" in output
+
+        for (f, s) in zip(opt.files, opt.sizes):
+            try:
+                output = u_boot_console.run_command(command + f)
+                assert str(s) in output
+            except:
+                assert False
+                opt.cleanup(build_dir)
+
+        # test symbolic link
         output = u_boot_console.run_command(command + "sym")
-        assert "100 bytes read in" in output
-    except:
-        sqfs_clean(cons)
-    sqfs_clean(cons)
+        assert str(opt.sizes[0]) in output
+
+        # remove generated files
+        opt.cleanup(build_dir)
diff --git a/test/py/tests/test_fs/test_squashfs/test_sqfs_ls.py b/test/py/tests/test_fs/test_squashfs/test_sqfs_ls.py
index 3a7b75c778..a0dca2e2fc 100644
--- a/test/py/tests/test_fs/test_squashfs/test_sqfs_ls.py
+++ b/test/py/tests/test_fs/test_squashfs/test_sqfs_ls.py
@@ -12,16 +12,25 @@ from sqfs_common import *
 @pytest.mark.buildconfigspec('fs_squashfs')
 @pytest.mark.requiredtool('mksquashfs')
 def test_sqfs_ls(u_boot_console):
-    cons = u_boot_console
-    sqfs_generate_image(cons)
-    path = os.path.join(cons.config.build_dir, "sqfs")
-    try:
+    build_dir = u_boot_console.config.build_dir
+    for opt in comp_opts:
+        try:
+            opt.gen_image(build_dir)
+        except RuntimeError:
+            opt.clean_source(build_dir)
+            # skip unsupported compression types
+            continue
+        path = os.path.join(build_dir, "sqfs-" + opt.name)
         output = u_boot_console.run_command("host bind 0 " + path)
-        output = u_boot_console.run_command("sqfsls host 0")
-        assert "4 file(s), 0 dir(s)" in output
-        assert "<SYM>   sym" in output
-        output = u_boot_console.run_command("sqfsls host 0 xxx")
-        assert "** Cannot find directory. **" in output
-    except:
-        sqfs_clean(cons)
-    sqfs_clean(cons)
+
+        try:
+            # list files in root directory
+            output = u_boot_console.run_command("sqfsls host 0")
+            assert str(len(opt.files) + 1) + " file(s), 0 dir(s)" in output
+            assert "<SYM>   sym" in output
+            output = u_boot_console.run_command("sqfsls host 0 xxx")
+            assert "** Cannot find directory. **" in output
+        except:
+            opt.cleanup(build_dir)
+            assert False
+        opt.cleanup(build_dir)