5 files changed, 217 insertions, 152 deletions

diff --git a/test/cmd_ut.c b/test/cmd_ut.c
index 14333423a1..6b24f463f3 100644
--- a/test/cmd_ut.c
+++ b/test/cmd_ut.c
@@ -8,9 +8,34 @@
 #include <common.h>
 #include <command.h>
 #include <test/suites.h>
+#include <test/test.h>
 
 static int do_ut_all(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[]);
 
+int cmd_ut_category(const char *name, struct unit_test *tests, int n_ents,
+		    int argc, char * const argv[])
+{
+	struct unit_test_state uts = { .fail_count = 0 };
+	struct unit_test *test;
+
+	if (argc == 1)
+		printf("Running %d %s tests\n", n_ents, name);
+
+	for (test = tests; test < tests + n_ents; test++) {
+		if (argc > 1 && strcmp(argv[1], test->name))
+			continue;
+		printf("Test: %s\n", test->name);
+
+		uts.start = mallinfo();
+
+		test->func(&uts);
+	}
+
+	printf("Failures: %d\n", uts.fail_count);
+
+	return uts.fail_count ? CMD_RET_FAILURE : 0;
+}
+
 static cmd_tbl_t cmd_ut_sub[] = {
 	U_BOOT_CMD_MKENT(all, CONFIG_SYS_MAXARGS, 1, do_ut_all, "", ""),
 #if defined(CONFIG_UT_DM)
@@ -25,6 +50,10 @@ static cmd_tbl_t cmd_ut_sub[] = {
 #ifdef CONFIG_UT_TIME
 	U_BOOT_CMD_MKENT(time, CONFIG_SYS_MAXARGS, 1, do_ut_time, "", ""),
 #endif
+#ifdef CONFIG_SANDBOX
+	U_BOOT_CMD_MKENT(compression, CONFIG_SYS_MAXARGS, 1, do_ut_compression,
+			 "", ""),
+#endif
 };
 
 static int do_ut_all(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
@@ -77,6 +106,9 @@ static char ut_help_text[] =
 #ifdef CONFIG_UT_TIME
 	"ut time - Very basic test of time functions\n"
 #endif
+#ifdef CONFIG_SANDBOX
+	"ut compression - Test compressors and bootm decompression\n"
+#endif
 	;
 #endif
 
diff --git a/test/compression.c b/test/compression.c
index be4e04e6cc..fe27ad66ea 100644
--- a/test/compression.c
+++ b/test/compression.c
@@ -4,8 +4,6 @@
  * SPDX-License-Identifier:	GPL-2.0+
  */
 
-#define DEBUG
-
 #include <common.h>
 #include <bootm.h>
 #include <command.h>
@@ -21,6 +19,9 @@
 #include <lzma/LzmaTools.h>
 
 #include <linux/lzo.h>
+#include <test/compression.h>
+#include <test/suites.h>
+#include <test/ut.h>
 
 static const char plain[] =
 	"I am a highly compressable bit of text.\n"
@@ -120,10 +121,11 @@ static const unsigned long lz4_compressed_size = 276;
 
 #define TEST_BUFFER_SIZE	512
 
-typedef int (*mutate_func)(void *, unsigned long, void *, unsigned long,
-			   unsigned long *);
+typedef int (*mutate_func)(struct unit_test_state *uts, void *, unsigned long,
+			   void *, unsigned long, unsigned long *);
 
-static int compress_using_gzip(void *in, unsigned long in_size,
+static int compress_using_gzip(struct unit_test_state *uts,
+			       void *in, unsigned long in_size,
 			       void *out, unsigned long out_max,
 			       unsigned long *out_size)
 {
@@ -137,7 +139,8 @@ static int compress_using_gzip(void *in, unsigned long in_size,
 	return ret;
 }
 
-static int uncompress_using_gzip(void *in, unsigned long in_size,
+static int uncompress_using_gzip(struct unit_test_state *uts,
+				 void *in, unsigned long in_size,
 				 void *out, unsigned long out_max,
 				 unsigned long *out_size)
 {
@@ -151,13 +154,14 @@ static int uncompress_using_gzip(void *in, unsigned long in_size,
 	return ret;
 }
 
-static int compress_using_bzip2(void *in, unsigned long in_size,
+static int compress_using_bzip2(struct unit_test_state *uts,
+				void *in, unsigned long in_size,
 				void *out, unsigned long out_max,
 				unsigned long *out_size)
 {
 	/* There is no bzip2 compression in u-boot, so fake it. */
-	assert(in_size == strlen(plain));
-	assert(memcmp(plain, in, in_size) == 0);
+	ut_asserteq(in_size, strlen(plain));
+	ut_asserteq(0, memcmp(plain, in, in_size));
 
 	if (bzip2_compressed_size > out_max)
 		return -1;
@@ -169,7 +173,8 @@ static int compress_using_bzip2(void *in, unsigned long in_size,
 	return 0;
 }
 
-static int uncompress_using_bzip2(void *in, unsigned long in_size,
+static int uncompress_using_bzip2(struct unit_test_state *uts,
+				  void *in, unsigned long in_size,
 				  void *out, unsigned long out_max,
 				  unsigned long *out_size)
 {
@@ -184,13 +189,14 @@ static int uncompress_using_bzip2(void *in, unsigned long in_size,
 	return (ret != BZ_OK);
 }
 
-static int compress_using_lzma(void *in, unsigned long in_size,
+static int compress_using_lzma(struct unit_test_state *uts,
+			       void *in, unsigned long in_size,
 			       void *out, unsigned long out_max,
 			       unsigned long *out_size)
 {
 	/* There is no lzma compression in u-boot, so fake it. */
-	assert(in_size == strlen(plain));
-	assert(memcmp(plain, in, in_size) == 0);
+	ut_asserteq(in_size,  strlen(plain));
+	ut_asserteq(0, memcmp(plain, in, in_size));
 
 	if (lzma_compressed_size > out_max)
 		return -1;
@@ -202,7 +208,8 @@ static int compress_using_lzma(void *in, unsigned long in_size,
 	return 0;
 }
 
-static int uncompress_using_lzma(void *in, unsigned long in_size,
+static int uncompress_using_lzma(struct unit_test_state *uts,
+				 void *in, unsigned long in_size,
 				 void *out, unsigned long out_max,
 				 unsigned long *out_size)
 {
@@ -216,13 +223,14 @@ static int uncompress_using_lzma(void *in, unsigned long in_size,
 	return (ret != SZ_OK);
 }
 
-static int compress_using_lzo(void *in, unsigned long in_size,
+static int compress_using_lzo(struct unit_test_state *uts,
+			      void *in, unsigned long in_size,
 			      void *out, unsigned long out_max,
 			      unsigned long *out_size)
 {
 	/* There is no lzo compression in u-boot, so fake it. */
-	assert(in_size == strlen(plain));
-	assert(memcmp(plain, in, in_size) == 0);
+	ut_asserteq(in_size,  strlen(plain));
+	ut_asserteq(0, memcmp(plain, in, in_size));
 
 	if (lzo_compressed_size > out_max)
 		return -1;
@@ -234,7 +242,8 @@ static int compress_using_lzo(void *in, unsigned long in_size,
 	return 0;
 }
 
-static int uncompress_using_lzo(void *in, unsigned long in_size,
+static int uncompress_using_lzo(struct unit_test_state *uts,
+				void *in, unsigned long in_size,
 				void *out, unsigned long out_max,
 				unsigned long *out_size)
 {
@@ -249,13 +258,14 @@ static int uncompress_using_lzo(void *in, unsigned long in_size,
 	return (ret != LZO_E_OK);
 }
 
-static int compress_using_lz4(void *in, unsigned long in_size,
+static int compress_using_lz4(struct unit_test_state *uts,
+			      void *in, unsigned long in_size,
 			      void *out, unsigned long out_max,
 			      unsigned long *out_size)
 {
 	/* There is no lz4 compression in u-boot, so fake it. */
-	assert(in_size == strlen(plain));
-	assert(memcmp(plain, in, in_size) == 0);
+	ut_asserteq(in_size,  strlen(plain));
+	ut_asserteq(0, memcmp(plain, in, in_size));
 
 	if (lz4_compressed_size > out_max)
 		return -1;
@@ -267,7 +277,8 @@ static int compress_using_lz4(void *in, unsigned long in_size,
 	return 0;
 }
 
-static int uncompress_using_lz4(void *in, unsigned long in_size,
+static int uncompress_using_lz4(struct unit_test_state *uts,
+				void *in, unsigned long in_size,
 				void *out, unsigned long out_max,
 				unsigned long *out_size)
 {
@@ -288,73 +299,69 @@ static int uncompress_using_lz4(void *in, unsigned long in_size,
 	goto out; \
 }
 
-static int run_test(char *name, mutate_func compress, mutate_func uncompress)
-{
-	ulong orig_size, compressed_size, uncompressed_size;
+struct buf_state {
+	ulong orig_size;
+	ulong compressed_size;
+	ulong uncompressed_size;
 	void *orig_buf;
-	void *compressed_buf = NULL;
-	void *uncompressed_buf = NULL;
-	void *compare_buf = NULL;
+	void *compressed_buf;
+	void *uncompressed_buf;
+	void *compare_buf;
+};
+
+static int run_test_internal(struct unit_test_state *uts, char *name,
+			     mutate_func compress, mutate_func uncompress,
+			     struct buf_state *buf)
+{
 	int ret;
 
-	printf(" testing %s ...\n", name);
-
-	orig_buf = (void *)plain;
-	orig_size = strlen(orig_buf); /* Trailing NULL not included. */
-	errcheck(orig_size > 0);
-
-	compressed_size = uncompressed_size = TEST_BUFFER_SIZE;
-	compressed_buf = malloc(compressed_size);
-	errcheck(compressed_buf != NULL);
-	uncompressed_buf = malloc(uncompressed_size);
-	errcheck(uncompressed_buf != NULL);
-	compare_buf = malloc(uncompressed_size);
-	errcheck(compare_buf != NULL);
-
 	/* Compress works as expected. */
-	printf("\torig_size:%lu\n", orig_size);
-	memset(compressed_buf, 'A', TEST_BUFFER_SIZE);
-	errcheck(compress(orig_buf, orig_size,
-			compressed_buf, compressed_size,
-			&compressed_size) == 0);
-	printf("\tcompressed_size:%lu\n", compressed_size);
-	errcheck(compressed_size > 0);
-	errcheck(compressed_size < orig_size);
-	errcheck(((char *)compressed_buf)[compressed_size-1] != 'A');
-	errcheck(((char *)compressed_buf)[compressed_size] == 'A');
+	printf("\torig_size:%lu\n", buf->orig_size);
+	memset(buf->compressed_buf, 'A', TEST_BUFFER_SIZE);
+	errcheck(compress(uts, buf->orig_buf, buf->orig_size,
+			  buf->compressed_buf, buf->compressed_size,
+			  &buf->compressed_size) == 0);
+	printf("\tcompressed_size:%lu\n", buf->compressed_size);
+	errcheck(buf->compressed_size > 0);
+	errcheck(buf->compressed_size < buf->orig_size);
+	errcheck(((char *)buf->compressed_buf)[buf->compressed_size - 1] !=
+			'A');
+	errcheck(((char *)buf->compressed_buf)[buf->compressed_size] == 'A');
 
 	/* Uncompresses with space remaining. */
-	errcheck(uncompress(compressed_buf, compressed_size,
-			  uncompressed_buf, uncompressed_size,
-			  &uncompressed_size) == 0);
-	printf("\tuncompressed_size:%lu\n", uncompressed_size);
-	errcheck(uncompressed_size == orig_size);
-	errcheck(memcmp(orig_buf, uncompressed_buf, orig_size) == 0);
+	errcheck(uncompress(uts, buf->compressed_buf, buf->compressed_size,
+			    buf->uncompressed_buf, buf->uncompressed_size,
+			    &buf->uncompressed_size) == 0);
+	printf("\tuncompressed_size:%lu\n", buf->uncompressed_size);
+	errcheck(buf->uncompressed_size == buf->orig_size);
+	errcheck(memcmp(buf->orig_buf, buf->uncompressed_buf,
+			buf->orig_size) == 0);
 
 	/* Uncompresses with exactly the right size output buffer. */
-	memset(uncompressed_buf, 'A', TEST_BUFFER_SIZE);
-	errcheck(uncompress(compressed_buf, compressed_size,
-			  uncompressed_buf, orig_size,
-			  &uncompressed_size) == 0);
-	errcheck(uncompressed_size == orig_size);
-	errcheck(memcmp(orig_buf, uncompressed_buf, orig_size) == 0);
-	errcheck(((char *)uncompressed_buf)[orig_size] == 'A');
+	memset(buf->uncompressed_buf, 'A', TEST_BUFFER_SIZE);
+	errcheck(uncompress(uts, buf->compressed_buf, buf->compressed_size,
+			    buf->uncompressed_buf, buf->orig_size,
+			    &buf->uncompressed_size) == 0);
+	errcheck(buf->uncompressed_size == buf->orig_size);
+	errcheck(memcmp(buf->orig_buf, buf->uncompressed_buf,
+			buf->orig_size) == 0);
+	errcheck(((char *)buf->uncompressed_buf)[buf->orig_size] == 'A');
 
 	/* Make sure compression does not over-run. */
-	memset(compare_buf, 'A', TEST_BUFFER_SIZE);
-	ret = compress(orig_buf, orig_size,
-		       compare_buf, compressed_size - 1,
+	memset(buf->compare_buf, 'A', TEST_BUFFER_SIZE);
+	ret = compress(uts, buf->orig_buf, buf->orig_size,
+		       buf->compare_buf, buf->compressed_size - 1,
 		       NULL);
-	errcheck(((char *)compare_buf)[compressed_size] == 'A');
+	errcheck(((char *)buf->compare_buf)[buf->compressed_size] == 'A');
 	errcheck(ret != 0);
 	printf("\tcompress does not overrun\n");
 
 	/* Make sure decompression does not over-run. */
-	memset(compare_buf, 'A', TEST_BUFFER_SIZE);
-	ret = uncompress(compressed_buf, compressed_size,
-			 compare_buf, uncompressed_size - 1,
+	memset(buf->compare_buf, 'A', TEST_BUFFER_SIZE);
+	ret = uncompress(uts, buf->compressed_buf, buf->compressed_size,
+			 buf->compare_buf, buf->uncompressed_size - 1,
 			 NULL);
-	errcheck(((char *)compare_buf)[uncompressed_size - 1] == 'A');
+	errcheck(((char *)buf->compare_buf)[buf->uncompressed_size - 1] == 'A');
 	errcheck(ret != 0);
 	printf("\tuncompress does not overrun\n");
 
@@ -362,32 +369,76 @@ static int run_test(char *name, mutate_func compress, mutate_func uncompress)
 	ret = 0;
 
 out:
+	return ret;
+}
+
+static int run_test(struct unit_test_state *uts, char *name,
+		    mutate_func compress, mutate_func uncompress)
+{
+	struct buf_state sbuf, *buf = &sbuf;
+	int ret;
+
+	printf(" testing %s ...\n", name);
+
+	buf->orig_buf = (void *)plain;
+	buf->orig_size = strlen(buf->orig_buf); /* Trailing NUL not included */
+	errcheck(buf->orig_size > 0);
+
+	buf->compressed_size = TEST_BUFFER_SIZE;
+	buf->uncompressed_size = TEST_BUFFER_SIZE;
+	buf->compressed_buf = malloc(buf->compressed_size);
+	errcheck(buf->compressed_buf);
+	buf->uncompressed_buf = malloc(buf->uncompressed_size);
+	errcheck(buf->uncompressed_buf);
+	buf->compare_buf = malloc(buf->uncompressed_size);
+	errcheck(buf->compare_buf);
+
+	ret = run_test_internal(uts, name, compress, uncompress, buf);
+out:
 	printf(" %s: %s\n", name, ret == 0 ? "ok" : "FAILED");
 
-	free(compare_buf);
-	free(uncompressed_buf);
-	free(compressed_buf);
+	free(buf->compare_buf);
+	free(buf->uncompressed_buf);
+	free(buf->compressed_buf);
 
 	return ret;
 }
 
-static int do_ut_compression(cmd_tbl_t *cmdtp, int flag, int argc,
-			     char *const argv[])
+static int compression_test_gzip(struct unit_test_state *uts)
 {
-	int err = 0;
+	return run_test(uts, "gzip", compress_using_gzip,
+			uncompress_using_gzip);
+}
+COMPRESSION_TEST(compression_test_gzip, 0);
+
+static int compression_test_bzip2(struct unit_test_state *uts)
+{
+	return run_test(uts, "bzip2", compress_using_bzip2,
+			uncompress_using_bzip2);
+}
+COMPRESSION_TEST(compression_test_bzip2, 0);
 
-	err += run_test("gzip", compress_using_gzip, uncompress_using_gzip);
-	err += run_test("bzip2", compress_using_bzip2, uncompress_using_bzip2);
-	err += run_test("lzma", compress_using_lzma, uncompress_using_lzma);
-	err += run_test("lzo", compress_using_lzo, uncompress_using_lzo);
-	err += run_test("lz4", compress_using_lz4, uncompress_using_lz4);
+static int compression_test_lzma(struct unit_test_state *uts)
+{
+	return run_test(uts, "lzma", compress_using_lzma,
+			uncompress_using_lzma);
+}
+COMPRESSION_TEST(compression_test_lzma, 0);
 
-	printf("ut_compression %s\n", err == 0 ? "ok" : "FAILED");
+static int compression_test_lzo(struct unit_test_state *uts)
+{
+	return run_test(uts, "lzo", compress_using_lzo, uncompress_using_lzo);
+}
+COMPRESSION_TEST(compression_test_lzo, 0);
 
-	return err;
+static int compression_test_lz4(struct unit_test_state *uts)
+{
+	return run_test(uts, "lz4", compress_using_lz4, uncompress_using_lz4);
 }
+COMPRESSION_TEST(compression_test_lz4, 0);
 
-static int compress_using_none(void *in, unsigned long in_size,
+static int compress_using_none(struct unit_test_state *uts,
+			       void *in, unsigned long in_size,
 			       void *out, unsigned long out_max,
 			       unsigned long *out_size)
 {
@@ -405,7 +456,8 @@ static int compress_using_none(void *in, unsigned long in_size,
  * @compress:	Our function to compress data
  * @return 0 if OK, non-zero on failure
  */
-static int run_bootm_test(int comp_type, mutate_func compress)
+static int run_bootm_test(struct unit_test_state *uts, int comp_type,
+			  mutate_func compress)
 {
 	ulong compress_size = 1024;
 	void *compress_buff;
@@ -418,20 +470,18 @@ static int run_bootm_test(int comp_type, mutate_func compress)
 	printf("Testing: %s\n", genimg_get_comp_name(comp_type));
 	compress_buff = map_sysmem(image_start, 0);
 	unc_len = strlen(plain);
-	compress((void *)plain, unc_len, compress_buff, compress_size,
+	compress(uts, (void *)plain, unc_len, compress_buff, compress_size,
 		 &compress_size);
 	err = bootm_decomp_image(comp_type, load_addr, image_start,
 				 IH_TYPE_KERNEL, map_sysmem(load_addr, 0),
 				 compress_buff, compress_size, unc_len,
 				 &load_end);
-	if (err)
-		return err;
+	ut_assertok(err);
 	err = bootm_decomp_image(comp_type, load_addr, image_start,
 				 IH_TYPE_KERNEL, map_sysmem(load_addr, 0),
 				 compress_buff, compress_size, unc_len - 1,
 				 &load_end);
-	if (!err)
-		return -EINVAL;
+	ut_assert(err);
 
 	/* We can't detect corruption when not decompressing */
 	if (comp_type == IH_COMP_NONE)
@@ -442,35 +492,52 @@ static int run_bootm_test(int comp_type, mutate_func compress)
 				 IH_TYPE_KERNEL, map_sysmem(load_addr, 0),
 				 compress_buff, compress_size, 0x10000,
 				 &load_end);
-	if (!err)
-		return -EINVAL;
+	ut_assert(err);
 
 	return 0;
 }
 
-static int do_ut_image_decomp(cmd_tbl_t *cmdtp, int flag, int argc,
-			      char *const argv[])
+static int compression_test_bootm_gzip(struct unit_test_state *uts)
 {
-	int err = 0;
+	return run_bootm_test(uts, IH_COMP_GZIP, compress_using_gzip);
+}
+COMPRESSION_TEST(compression_test_bootm_gzip, 0);
+
+static int compression_test_bootm_bzip2(struct unit_test_state *uts)
+{
+	return run_bootm_test(uts, IH_COMP_BZIP2, compress_using_bzip2);
+}
+COMPRESSION_TEST(compression_test_bootm_bzip2, 0);
 
-	err = run_bootm_test(IH_COMP_GZIP, compress_using_gzip);
-	err |= run_bootm_test(IH_COMP_BZIP2, compress_using_bzip2);
-	err |= run_bootm_test(IH_COMP_LZMA, compress_using_lzma);
-	err |= run_bootm_test(IH_COMP_LZO, compress_using_lzo);
-	err |= run_bootm_test(IH_COMP_LZ4, compress_using_lz4);
-	err |= run_bootm_test(IH_COMP_NONE, compress_using_none);
+static int compression_test_bootm_lzma(struct unit_test_state *uts)
+{
+	return run_bootm_test(uts, IH_COMP_LZMA, compress_using_lzma);
+}
+COMPRESSION_TEST(compression_test_bootm_lzma, 0);
 
-	printf("ut_image_decomp %s\n", err == 0 ? "ok" : "FAILED");
+static int compression_test_bootm_lzo(struct unit_test_state *uts)
+{
+	return run_bootm_test(uts, IH_COMP_LZO, compress_using_lzo);
+}
+COMPRESSION_TEST(compression_test_bootm_lzo, 0);
 
-	return 0;
+static int compression_test_bootm_lz4(struct unit_test_state *uts)
+{
+	return run_bootm_test(uts, IH_COMP_LZ4, compress_using_lz4);
 }
+COMPRESSION_TEST(compression_test_bootm_lz4, 0);
 
-U_BOOT_CMD(
-	ut_compression,	5,	1,	do_ut_compression,
-	"Basic test of compressors: gzip bzip2 lzma lzo", ""
-);
+static int compression_test_bootm_none(struct unit_test_state *uts)
+{
+	return run_bootm_test(uts, IH_COMP_NONE, compress_using_none);
+}
+COMPRESSION_TEST(compression_test_bootm_none, 0);
 
-U_BOOT_CMD(
-	ut_image_decomp,	5,	1, do_ut_image_decomp,
-	"Basic test of bootm decompression", ""
-);
+int do_ut_compression(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
+{
+	struct unit_test *tests = ll_entry_start(struct unit_test,
+						 compression_test);
+	const int n_ents = ll_entry_count(struct unit_test, compression_test);
+
+	return cmd_ut_category("compression", tests, n_ents, argc, argv);
+}
diff --git a/test/env/cmd_ut_env.c b/test/env/cmd_ut_env.c
index 893e5e6a6d..096afa83dd 100644
--- a/test/env/cmd_ut_env.c
+++ b/test/env/cmd_ut_env.c
@@ -15,23 +15,6 @@ int do_ut_env(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
 {
 	struct unit_test *tests = ll_entry_start(struct unit_test, env_test);
 	const int n_ents = ll_entry_count(struct unit_test, env_test);
-	struct unit_test_state uts = { .fail_count = 0 };
-	struct unit_test *test;
 
-	if (argc == 1)
-		printf("Running %d environment tests\n", n_ents);
-
-	for (test = tests; test < tests + n_ents; test++) {
-		if (argc > 1 && strcmp(argv[1], test->name))
-			continue;
-		printf("Test: %s\n", test->name);
-
-		uts.start = mallinfo();
-
-		test->func(&uts);
-	}
-
-	printf("Failures: %d\n", uts.fail_count);
-
-	return uts.fail_count ? CMD_RET_FAILURE : 0;
+	return cmd_ut_category("environment", tests, n_ents, argc, argv);
 }
diff --git a/test/overlay/cmd_ut_overlay.c b/test/overlay/cmd_ut_overlay.c
index c730a11f51..6279e6d0c2 100644
--- a/test/overlay/cmd_ut_overlay.c
+++ b/test/overlay/cmd_ut_overlay.c
@@ -14,6 +14,7 @@
 
 #include <test/ut.h>
 #include <test/overlay.h>
+#include <test/suites.h>
 
 /* 4k ought to be enough for anybody */
 #define FDT_COPY_SIZE	(4 * SZ_1K)
@@ -221,7 +222,6 @@ int do_ut_overlay(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
 						 overlay_test);
 	const int n_ents = ll_entry_count(struct unit_test, overlay_test);
 	struct unit_test_state *uts;
-	struct unit_test *test;
 	void *fdt_base = &__dtb_test_fdt_base_begin;
 	void *fdt_overlay = &__dtb_test_fdt_overlay_begin;
 	void *fdt_overlay_stacked = &__dtb_test_fdt_overlay_stacked_begin;
@@ -280,24 +280,7 @@ int do_ut_overlay(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
 	/* Apply the stacked overlay */
 	ut_assertok(fdt_overlay_apply(fdt_base_copy, fdt_overlay_stacked_copy));
 
-	if (argc == 1)
-		printf("Running %d environment tests\n", n_ents);
-
-	for (test = tests; test < tests + n_ents; test++) {
-		if (argc > 1 && strcmp(argv[1], test->name))
-			continue;
-		printf("Test: %s\n", test->name);
-
-		uts->start = mallinfo();
-
-		test->func(uts);
-	}
-
-	printf("Failures: %d\n", uts->fail_count);
-	if (!uts->fail_count)
-		ret = 0;
-	else
-		ret = CMD_RET_FAILURE;
+	ret = cmd_ut_category("overlay", tests, n_ents, argc, argv);
 
 	free(fdt_overlay_stacked_copy);
 err3:
diff --git a/test/py/conftest.py b/test/py/conftest.py
index 6e66a48c15..3fe91e8746 100644
--- a/test/py/conftest.py
+++ b/test/py/conftest.py
@@ -200,7 +200,7 @@ def pytest_configure(config):
         import u_boot_console_exec_attach
         console = u_boot_console_exec_attach.ConsoleExecAttach(log, ubconfig)
 
-re_ut_test_list = re.compile(r'_u_boot_list_2_(dm|env)_test_2_\1_test_(.*)\s*$')
+re_ut_test_list = re.compile(r'_u_boot_list_2_(.*)_test_2_\1_test_(.*)\s*$')
 def generate_ut_subtest(metafunc, fixture_name):
     """Provide parametrization for a ut_subtest fixture.