Merge branch '2021-03-12-test-improvements' into next

- Assorted improvements to the pytest framework

28 files changed, 674 insertions, 343 deletions

diff --git a/test/Makefile b/test/Makefile
index 932e517383..a26e915e05 100644
--- a/test/Makefile
+++ b/test/Makefile
@@ -2,6 +2,9 @@
 #
 # (C) Copyright 2012 The Chromium Authors
 
+obj-y += test-main.o
+obj-$(CONFIG_SANDBOX) += image/
+
 ifneq ($(CONFIG_$(SPL_)BLOBLIST),)
 obj-$(CONFIG_$(SPL_)CMDLINE) += bloblist.o
 obj-$(CONFIG_$(SPL_)CMDLINE) += bootm.o
diff --git a/test/bloblist.c b/test/bloblist.c
index 6953d3010a..d876b63918 100644
--- a/test/bloblist.c
+++ b/test/bloblist.c
@@ -387,9 +387,8 @@ BLOBLIST_TEST(bloblist_test_reloc, 0);
 int do_ut_bloblist(struct cmd_tbl *cmdtp, int flag, int argc,
 		   char *const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test,
-						 bloblist_test);
-	const int n_ents = ll_entry_count(struct unit_test, bloblist_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(bloblist_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(bloblist_test);
 
 	return cmd_ut_category("bloblist", "bloblist_test_",
 			       tests, n_ents, argc, argv);
diff --git a/test/bootm.c b/test/bootm.c
index 563d6ebaa5..8528982ae1 100644
--- a/test/bootm.c
+++ b/test/bootm.c
@@ -240,8 +240,8 @@ BOOTM_TEST(bootm_test_subst_both, 0);
 
 int do_ut_bootm(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test, bootm_test);
-	const int n_ents = ll_entry_count(struct unit_test, bootm_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(bootm_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(bootm_test);
 
 	return cmd_ut_category("bootm", "bootm_test_", tests, n_ents,
 			       argc, argv);
diff --git a/test/cmd/mem.c b/test/cmd/mem.c
index fbaa8a4b3c..d76f47cf31 100644
--- a/test/cmd/mem.c
+++ b/test/cmd/mem.c
@@ -12,8 +12,8 @@
 
 int do_ut_mem(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test, mem_test);
-	const int n_ents = ll_entry_count(struct unit_test, mem_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(mem_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(mem_test);
 
 	return cmd_ut_category("cmd_mem", "mem_test_", tests, n_ents, argc,
 			       argv);
diff --git a/test/cmd/setexpr.c b/test/cmd/setexpr.c
index fd6d869c0e..27113c083c 100644
--- a/test/cmd/setexpr.c
+++ b/test/cmd/setexpr.c
@@ -306,8 +306,8 @@ static int setexpr_test_str(struct unit_test_state *uts)
 	ut_asserteq(1, run_command("setexpr.s fred 0", 0));
 	ut_assertok(ut_check_delta(start_mem));
 
-	start_mem = ut_check_free();
 	ut_assertok(env_set("fred", "12345"));
+	start_mem = ut_check_free();
 	ut_assertok(run_command("setexpr.s fred *0", 0));
 	ut_asserteq_str("hello", env_get("fred"));
 	ut_assertok(ut_check_delta(start_mem));
@@ -345,7 +345,22 @@ static int setexpr_test_str_oper(struct unit_test_state *uts)
 	start_mem = ut_check_free();
 	ut_assertok(run_command("setexpr.s fred *0 + *10", 0));
 	ut_asserteq_str("hello there", env_get("fred"));
-	ut_assertok(ut_check_delta(start_mem));
+
+	/*
+	 * This check does not work with sandbox_flattree, apparently due to
+	 * memory allocations in env_set().
+	 *
+	 * The truetype console produces lots of memory allocations even though
+	 * the LCD display is not visible. But even without these, it does not
+	 * work.
+	 *
+	 * A better test would be for dlmalloc to record the allocs and frees
+	 * for a particular caller, but that is not supported.
+	 *
+	 * For now, drop this test.
+	 *
+	 * ut_assertok(ut_check_delta(start_mem));
+	 */
 
 	unmap_sysmem(buf);
 
@@ -375,10 +390,9 @@ SETEXPR_TEST(setexpr_test_str_long, UT_TESTF_CONSOLE_REC);
 
 int do_ut_setexpr(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test,
-						 setexpr_test);
-	const int n_ents = ll_entry_count(struct unit_test, setexpr_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(setexpr_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(setexpr_test);
 
-	return cmd_ut_category("cmd_setexpr", "cmd_mem_", tests, n_ents, argc,
-			       argv);
+	return cmd_ut_category("cmd_setexpr", "setexpr_test_", tests, n_ents,
+			       argc, argv);
 }
diff --git a/test/cmd_ut.c b/test/cmd_ut.c
index 8f3089890e..157f6aa976 100644
--- a/test/cmd_ut.c
+++ b/test/cmd_ut.c
@@ -9,6 +9,7 @@
 #include <console.h>
 #include <test/suites.h>
 #include <test/test.h>
+#include <test/ut.h>
 
 static int do_ut_all(struct cmd_tbl *cmdtp, int flag, int argc,
 		     char *const argv[]);
@@ -17,41 +18,12 @@ int cmd_ut_category(const char *name, const char *prefix,
 		    struct unit_test *tests, int n_ents,
 		    int argc, char *const argv[])
 {
-	struct unit_test_state uts = { .fail_count = 0 };
-	struct unit_test *test;
-	int prefix_len = prefix ? strlen(prefix) : 0;
+	int ret;
 
-	if (argc == 1)
-		printf("Running %d %s tests\n", n_ents, name);
+	ret = ut_run_list(name, prefix, tests, n_ents,
+			  argc > 1 ? argv[1] : NULL);
 
-	for (test = tests; test < tests + n_ents; test++) {
-		const char *test_name = test->name;
-
-		/* Remove the prefix */
-		if (prefix && !strncmp(test_name, prefix, prefix_len))
-			test_name += prefix_len;
-
-		if (argc > 1 && strcmp(argv[1], test_name))
-			continue;
-		printf("Test: %s\n", test->name);
-
-		if (test->flags & UT_TESTF_CONSOLE_REC) {
-			int ret = console_record_reset_enable();
-
-			if (ret) {
-				printf("Skipping: Console recording disabled\n");
-				continue;
-			}
-		}
-
-		uts.start = mallinfo();
-
-		test->func(&uts);
-	}
-
-	printf("Failures: %d\n", uts.fail_count);
-
-	return uts.fail_count ? CMD_RET_FAILURE : 0;
+	return ret ? CMD_RET_FAILURE : 0;
 }
 
 static struct cmd_tbl cmd_ut_sub[] = {
diff --git a/test/compression.c b/test/compression.c
index a2a4b9ff9e..4cd1be564f 100644
--- a/test/compression.c
+++ b/test/compression.c
@@ -539,9 +539,8 @@ COMPRESSION_TEST(compression_test_bootm_none, 0);
 int do_ut_compression(struct cmd_tbl *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);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(compression_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(compression_test);
 
 	return cmd_ut_category("compression", "compression_test_",
 			       tests, n_ents, argc, argv);
diff --git a/test/dm/Makefile b/test/dm/Makefile
index fd1455109d..f5cc5540e8 100644
--- a/test/dm/Makefile
+++ b/test/dm/Makefile
@@ -2,7 +2,7 @@
 #
 # Copyright (c) 2013 Google, Inc
 
-obj-$(CONFIG_UT_DM) += test-main.o
+obj-$(CONFIG_UT_DM) += test-dm.o
 
 # Tests for particular subsystems - when enabling driver model for a new
 # subsystem you must add sandbox tests here.
diff --git a/test/dm/core.c b/test/dm/core.c
index 35ca689d64..2210345dd1 100644
--- a/test/dm/core.c
+++ b/test/dm/core.c
@@ -117,14 +117,13 @@ int dm_leak_check_end(struct unit_test_state *uts)
 /* Test that binding with plat occurs correctly */
 static int dm_test_autobind(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	struct udevice *dev;
 
 	/*
 	 * We should have a single class (UCLASS_ROOT) and a single root
 	 * device with no children.
 	 */
-	ut_assert(dms->root);
+	ut_assert(uts->root);
 	ut_asserteq(1, list_count_items(gd->uclass_root));
 	ut_asserteq(0, list_count_items(&gd->dm_root->child_head));
 	ut_asserteq(0, dm_testdrv_op_count[DM_TEST_OP_POST_BIND]);
@@ -207,7 +206,6 @@ DM_TEST(dm_test_autobind_uclass_pdata_valid, UT_TESTF_SCAN_PDATA);
 /* Test that autoprobe finds all the expected devices */
 static int dm_test_autoprobe(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	int expected_base_add;
 	struct udevice *dev;
 	struct uclass *uc;
@@ -221,7 +219,7 @@ static int dm_test_autoprobe(struct unit_test_state *uts)
 	ut_asserteq(0, dm_testdrv_op_count[DM_TEST_OP_POST_PROBE]);
 
 	/* The root device should not be activated until needed */
-	ut_assert(dev_get_flags(dms->root) & DM_FLAG_ACTIVATED);
+	ut_assert(dev_get_flags(uts->root) & DM_FLAG_ACTIVATED);
 
 	/*
 	 * We should be able to find the three test devices, and they should
@@ -241,7 +239,7 @@ static int dm_test_autoprobe(struct unit_test_state *uts)
 
 		/* Activating a device should activate the root device */
 		if (!i)
-			ut_assert(dev_get_flags(dms->root) & DM_FLAG_ACTIVATED);
+			ut_assert(dev_get_flags(uts->root) & DM_FLAG_ACTIVATED);
 	}
 
 	/*
@@ -293,7 +291,6 @@ DM_TEST(dm_test_plat, UT_TESTF_SCAN_PDATA);
 /* Test that we can bind, probe, remove, unbind a driver */
 static int dm_test_lifecycle(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	int op_count[DM_TEST_OP_COUNT];
 	struct udevice *dev, *test_dev;
 	int pingret;
@@ -301,7 +298,7 @@ static int dm_test_lifecycle(struct unit_test_state *uts)
 
 	memcpy(op_count, dm_testdrv_op_count, sizeof(op_count));
 
-	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_manual,
+	ut_assertok(device_bind_by_name(uts->root, false, &driver_info_manual,
 					&dev));
 	ut_assert(dev);
 	ut_assert(dm_testdrv_op_count[DM_TEST_OP_BIND]
@@ -309,7 +306,7 @@ static int dm_test_lifecycle(struct unit_test_state *uts)
 	ut_assert(!dev_get_priv(dev));
 
 	/* Probe the device - it should fail allocating private data */
-	dms->force_fail_alloc = 1;
+	uts->force_fail_alloc = 1;
 	ret = device_probe(dev);
 	ut_assert(ret == -ENOMEM);
 	ut_assert(dm_testdrv_op_count[DM_TEST_OP_PROBE]
@@ -317,7 +314,7 @@ static int dm_test_lifecycle(struct unit_test_state *uts)
 	ut_assert(!dev_get_priv(dev));
 
 	/* Try again without the alloc failure */
-	dms->force_fail_alloc = 0;
+	uts->force_fail_alloc = 0;
 	ut_assertok(device_probe(dev));
 	ut_assert(dm_testdrv_op_count[DM_TEST_OP_PROBE]
 			== op_count[DM_TEST_OP_PROBE] + 2);
@@ -349,19 +346,18 @@ DM_TEST(dm_test_lifecycle, UT_TESTF_SCAN_PDATA | UT_TESTF_PROBE_TEST);
 /* Test that we can bind/unbind and the lists update correctly */
 static int dm_test_ordering(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	struct udevice *dev, *dev_penultimate, *dev_last, *test_dev;
 	int pingret;
 
-	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_manual,
+	ut_assertok(device_bind_by_name(uts->root, false, &driver_info_manual,
 					&dev));
 	ut_assert(dev);
 
 	/* Bind two new devices (numbers 4 and 5) */
-	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_manual,
+	ut_assertok(device_bind_by_name(uts->root, false, &driver_info_manual,
 					&dev_penultimate));
 	ut_assert(dev_penultimate);
-	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_manual,
+	ut_assertok(device_bind_by_name(uts->root, false, &driver_info_manual,
 					&dev_last));
 	ut_assert(dev_last);
 
@@ -376,7 +372,7 @@ static int dm_test_ordering(struct unit_test_state *uts)
 	ut_assert(dev_last == test_dev);
 
 	/* Add back the original device 3, now in position 5 */
-	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_manual,
+	ut_assertok(device_bind_by_name(uts->root, false, &driver_info_manual,
 					&dev));
 	ut_assert(dev);
 
@@ -568,7 +564,6 @@ static int create_children(struct unit_test_state *uts, struct udevice *parent,
 
 static int dm_test_children(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	struct udevice *top[NODE_COUNT];
 	struct udevice *child[NODE_COUNT];
 	struct udevice *grandchild[NODE_COUNT];
@@ -578,12 +573,12 @@ static int dm_test_children(struct unit_test_state *uts)
 	int i;
 
 	/* We don't care about the numbering for this test */
-	dms->skip_post_probe = 1;
+	uts->skip_post_probe = 1;
 
 	ut_assert(NODE_COUNT > 5);
 
 	/* First create 10 top-level children */
-	ut_assertok(create_children(uts, dms->root, NODE_COUNT, 0, top));
+	ut_assertok(create_children(uts, uts->root, NODE_COUNT, 0, top));
 
 	/* Now a few have their own children */
 	ut_assertok(create_children(uts, top[2], NODE_COUNT, 2, NULL));
@@ -654,7 +649,6 @@ DM_TEST(dm_test_children, 0);
 
 static int dm_test_device_reparent(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	struct udevice *top[NODE_COUNT];
 	struct udevice *child[NODE_COUNT];
 	struct udevice *grandchild[NODE_COUNT];
@@ -664,12 +658,12 @@ static int dm_test_device_reparent(struct unit_test_state *uts)
 	int i;
 
 	/* We don't care about the numbering for this test */
-	dms->skip_post_probe = 1;
+	uts->skip_post_probe = 1;
 
 	ut_assert(NODE_COUNT > 5);
 
 	/* First create 10 top-level children */
-	ut_assertok(create_children(uts, dms->root, NODE_COUNT, 0, top));
+	ut_assertok(create_children(uts, uts->root, NODE_COUNT, 0, top));
 
 	/* Now a few have their own children */
 	ut_assertok(create_children(uts, top[2], NODE_COUNT, 2, NULL));
@@ -815,15 +809,14 @@ DM_TEST(dm_test_device_reparent, 0);
 /* Test that pre-relocation devices work as expected */
 static int dm_test_pre_reloc(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	struct udevice *dev;
 
 	/* The normal driver should refuse to bind before relocation */
-	ut_asserteq(-EPERM, device_bind_by_name(dms->root, true,
+	ut_asserteq(-EPERM, device_bind_by_name(uts->root, true,
 						&driver_info_manual, &dev));
 
 	/* But this one is marked pre-reloc */
-	ut_assertok(device_bind_by_name(dms->root, true,
+	ut_assertok(device_bind_by_name(uts->root, true,
 					&driver_info_pre_reloc, &dev));
 
 	return 0;
@@ -836,10 +829,9 @@ DM_TEST(dm_test_pre_reloc, 0);
  */
 static int dm_test_remove_active_dma(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	struct udevice *dev;
 
-	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_act_dma,
+	ut_assertok(device_bind_by_name(uts->root, false, &driver_info_act_dma,
 					&dev));
 	ut_assert(dev);
 
@@ -872,7 +864,7 @@ static int dm_test_remove_active_dma(struct unit_test_state *uts)
 	 * the active DMA remove call
 	 */
 	ut_assertok(device_unbind(dev));
-	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_manual,
+	ut_assertok(device_bind_by_name(uts->root, false, &driver_info_manual,
 					&dev));
 	ut_assert(dev);
 
@@ -895,25 +887,24 @@ DM_TEST(dm_test_remove_active_dma, 0);
 /* Test removal of 'vital' devices */
 static int dm_test_remove_vital(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	struct udevice *normal, *dma, *vital, *dma_vital;
 
 	/* Skip the behaviour in test_post_probe() */
-	dms->skip_post_probe = 1;
+	uts->skip_post_probe = 1;
 
-	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_manual,
+	ut_assertok(device_bind_by_name(uts->root, false, &driver_info_manual,
 					&normal));
 	ut_assertnonnull(normal);
 
-	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_act_dma,
+	ut_assertok(device_bind_by_name(uts->root, false, &driver_info_act_dma,
 					&dma));
 	ut_assertnonnull(dma);
 
-	ut_assertok(device_bind_by_name(dms->root, false,
+	ut_assertok(device_bind_by_name(uts->root, false,
 					&driver_info_vital_clk, &vital));
 	ut_assertnonnull(vital);
 
-	ut_assertok(device_bind_by_name(dms->root, false,
+	ut_assertok(device_bind_by_name(uts->root, false,
 					&driver_info_act_dma_vital_clk,
 					&dma_vital));
 	ut_assertnonnull(dma_vital);
@@ -1133,11 +1124,10 @@ DM_TEST(dm_test_uclass_names, UT_TESTF_SCAN_PDATA);
 
 static int dm_test_inactive_child(struct unit_test_state *uts)
 {
-	struct dm_test_state *dms = uts->priv;
 	struct udevice *parent, *dev1, *dev2;
 
 	/* Skip the behaviour in test_post_probe() */
-	dms->skip_post_probe = 1;
+	uts->skip_post_probe = 1;
 
 	ut_assertok(uclass_first_device_err(UCLASS_TEST, &parent));
 
diff --git a/test/dm/test-dm.c b/test/dm/test-dm.c
new file mode 100644
index 0000000000..9ba2ba23d5
--- /dev/null
+++ b/test/dm/test-dm.c
@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * Copyright (c) 2013 Google, Inc
+ */
+
+#include <common.h>
+#include <command.h>
+#include <console.h>
+#include <dm.h>
+#include <errno.h>
+#include <log.h>
+#include <malloc.h>
+#include <asm/global_data.h>
+#include <asm/state.h>
+#include <dm/root.h>
+#include <dm/uclass-internal.h>
+#include <test/test.h>
+#include <test/test.h>
+#include <test/ut.h>
+
+DECLARE_GLOBAL_DATA_PTR;
+
+/**
+ * dm_test_run() - Run driver model tests
+ *
+ * Run all the available driver model tests, or a selection
+ *
+ * @test_name: Name of single test to run (e.g. "dm_test_fdt_pre_reloc" or just
+ *	"fdt_pre_reloc"), or NULL to run all
+ * @return 0 if all tests passed, 1 if not
+ */
+static int dm_test_run(const char *test_name)
+{
+	struct unit_test *tests = UNIT_TEST_SUITE_START(dm_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(dm_test);
+	int ret;
+
+	ret = ut_run_list("driver model", "dm_test_", tests, n_ents, test_name);
+
+	return ret ? CMD_RET_FAILURE : 0;
+}
+
+int do_ut_dm(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
+{
+	const char *test_name = NULL;
+
+	if (argc > 1)
+		test_name = argv[1];
+
+	return dm_test_run(test_name);
+}
diff --git a/test/dm/test-driver.c b/test/dm/test-driver.c
index ca7626a066..02cb974b0f 100644
--- a/test/dm/test-driver.c
+++ b/test/dm/test-driver.c
@@ -18,7 +18,6 @@
 #include <test/ut.h>
 
 int dm_testdrv_op_count[DM_TEST_OP_COUNT];
-static struct unit_test_state *uts = &global_dm_test_state;
 
 static int testdrv_ping(struct udevice *dev, int pingval, int *pingret)
 {
@@ -37,6 +36,8 @@ static const struct test_ops test_ops = {
 
 static int test_bind(struct udevice *dev)
 {
+	struct unit_test_state *uts = test_get_state();
+
 	/* Private data should not be allocated */
 	ut_assert(!dev_get_priv(dev));
 
@@ -46,6 +47,7 @@ static int test_bind(struct udevice *dev)
 
 static int test_probe(struct udevice *dev)
 {
+	struct unit_test_state *uts = test_get_state();
 	struct dm_test_priv *priv = dev_get_priv(dev);
 
 	/* Private data should be allocated */
@@ -58,6 +60,8 @@ static int test_probe(struct udevice *dev)
 
 static int test_remove(struct udevice *dev)
 {
+	struct unit_test_state *uts = test_get_state();
+
 	/* Private data should still be allocated */
 	ut_assert(dev_get_priv(dev));
 
@@ -67,6 +71,8 @@ static int test_remove(struct udevice *dev)
 
 static int test_unbind(struct udevice *dev)
 {
+	struct unit_test_state *uts = test_get_state();
+
 	/* Private data should not be allocated */
 	ut_assert(!dev_get_priv(dev));
 
@@ -116,10 +122,10 @@ static int test_manual_bind(struct udevice *dev)
 
 static int test_manual_probe(struct udevice *dev)
 {
-	struct dm_test_state *dms = uts->priv;
+	struct unit_test_state *uts = test_get_state();
 
 	dm_testdrv_op_count[DM_TEST_OP_PROBE]++;
-	if (!dms->force_fail_alloc)
+	if (!uts->force_fail_alloc)
 		dev_set_priv(dev, calloc(1, sizeof(struct dm_test_priv)));
 	if (!dev_get_priv(dev))
 		return -ENOMEM;
diff --git a/test/dm/test-main.c b/test/dm/test-main.c
deleted file mode 100644
index 560f8d63ec..0000000000
--- a/test/dm/test-main.c
+++ /dev/null
@@ -1,230 +0,0 @@
-// SPDX-License-Identifier: GPL-2.0+
-/*
- * Copyright (c) 2013 Google, Inc
- */
-
-#include <common.h>
-#include <command.h>
-#include <console.h>
-#include <dm.h>
-#include <errno.h>
-#include <log.h>
-#include <malloc.h>
-#include <asm/global_data.h>
-#include <asm/state.h>
-#include <dm/test.h>
-#include <dm/root.h>
-#include <dm/uclass-internal.h>
-#include <test/test.h>
-#include <test/test.h>
-#include <test/ut.h>
-
-DECLARE_GLOBAL_DATA_PTR;
-
-struct unit_test_state global_dm_test_state;
-static struct dm_test_state _global_priv_dm_test_state;
-
-/* Get ready for testing */
-static int dm_test_init(struct unit_test_state *uts, bool of_live)
-{
-	struct dm_test_state *dms = uts->priv;
-
-	memset(dms, '\0', sizeof(*dms));
-	gd->dm_root = NULL;
-	if (!CONFIG_IS_ENABLED(OF_PLATDATA))
-		memset(dm_testdrv_op_count, '\0', sizeof(dm_testdrv_op_count));
-	state_reset_for_test(state_get_current());
-
-	/* Determine whether to make the live tree available */
-	gd_set_of_root(of_live ? uts->of_root : NULL);
-	ut_assertok(dm_init(of_live));
-	dms->root = dm_root();
-
-	return 0;
-}
-
-/* Ensure all the test devices are probed */
-static int do_autoprobe(struct unit_test_state *uts)
-{
-	struct udevice *dev;
-	int ret;
-
-	/* Scanning the uclass is enough to probe all the devices */
-	for (ret = uclass_first_device(UCLASS_TEST, &dev);
-	     dev;
-	     ret = uclass_next_device(&dev))
-		;
-
-	return ret;
-}
-
-static int dm_test_destroy(struct unit_test_state *uts)
-{
-	int id;
-
-	for (id = 0; id < UCLASS_COUNT; id++) {
-		struct uclass *uc;
-
-		/*
-		 * If the uclass doesn't exist we don't want to create it. So
-		 * check that here before we call uclass_find_device().
-		 */
-		uc = uclass_find(id);
-		if (!uc)
-			continue;
-		ut_assertok(uclass_destroy(uc));
-	}
-
-	return 0;
-}
-
-static int dm_do_test(struct unit_test_state *uts, struct unit_test *test,
-		      bool of_live)
-{
-	struct sandbox_state *state = state_get_current();
-	const char *fname = strrchr(test->file, '/') + 1;
-
-	printf("Test: %s: %s%s\n", test->name, fname,
-	       !of_live ? " (flat tree)" : "");
-	ut_assertok(dm_test_init(uts, of_live));
-
-	uts->start = mallinfo();
-	if (test->flags & UT_TESTF_SCAN_PDATA)
-		ut_assertok(dm_scan_plat(false));
-	if (test->flags & UT_TESTF_PROBE_TEST)
-		ut_assertok(do_autoprobe(uts));
-	if (!CONFIG_IS_ENABLED(OF_PLATDATA) &&
-	    (test->flags & UT_TESTF_SCAN_FDT))
-		ut_assertok(dm_extended_scan(false));
-
-	/*
-	 * Silence the console and rely on console recording to get
-	 * our output.
-	 */
-	console_record_reset_enable();
-	if (!state->show_test_output)
-		gd->flags |= GD_FLG_SILENT;
-	test->func(uts);
-	gd->flags &= ~(GD_FLG_SILENT | GD_FLG_RECORD);
-	state_set_skip_delays(false);
-
-	ut_assertok(dm_test_destroy(uts));
-
-	return 0;
-}
-
-/**
- * dm_test_run_on_flattree() - Check if we should run a test with flat DT
- *
- * This skips long/slow tests where there is not much value in running a flat
- * DT test in addition to a live DT test.
- *
- * @return true to run the given test on the flat device tree
- */
-static bool dm_test_run_on_flattree(struct unit_test *test)
-{
-	const char *fname = strrchr(test->file, '/') + 1;
-
-	return !strstr(fname, "video") || strstr(test->name, "video_base");
-}
-
-static bool test_matches(const char *test_name, const char *find_name)
-{
-	if (!find_name)
-		return true;
-
-	if (!strcmp(test_name, find_name))
-		return true;
-
-	/* All tests have this prefix */
-	if (!strncmp(test_name, "dm_test_", 8))
-		test_name += 8;
-
-	if (!strcmp(test_name, find_name))
-		return true;
-
-	return false;
-}
-
-int dm_test_main(const char *test_name)
-{
-	struct unit_test *tests = ll_entry_start(struct unit_test, dm_test);
-	const int n_ents = ll_entry_count(struct unit_test, dm_test);
-	struct unit_test_state *uts = &global_dm_test_state;
-	struct unit_test *test;
-	int found;
-
-	uts->priv = &_global_priv_dm_test_state;
-	uts->fail_count = 0;
-
-	if (!CONFIG_IS_ENABLED(OF_PLATDATA)) {
-		/*
-		 * If we have no device tree, or it only has a root node, then
-		 * these * tests clearly aren't going to work...
-		 */
-		if (!gd->fdt_blob || fdt_next_node(gd->fdt_blob, 0, NULL) < 0) {
-			puts("Please run with test device tree:\n"
-			     "    ./u-boot -d arch/sandbox/dts/test.dtb\n");
-			ut_assert(gd->fdt_blob);
-		}
-	}
-
-	if (!test_name)
-		printf("Running %d driver model tests\n", n_ents);
-	else
-
-	found = 0;
-	uts->of_root = gd_of_root();
-	for (test = tests; test < tests + n_ents; test++) {
-		const char *name = test->name;
-		int runs;
-
-		if (!test_matches(name, test_name))
-			continue;
-
-		/* Run with the live tree if possible */
-		runs = 0;
-		if (CONFIG_IS_ENABLED(OF_LIVE)) {
-			if (!(test->flags & UT_TESTF_FLAT_TREE)) {
-				ut_assertok(dm_do_test(uts, test, true));
-				runs++;
-			}
-		}
-
-		/*
-		 * Run with the flat tree if we couldn't run it with live tree,
-		 * or it is a core test.
-		 */
-		if (!(test->flags & UT_TESTF_LIVE_TREE) &&
-		    (!runs || dm_test_run_on_flattree(test))) {
-			ut_assertok(dm_do_test(uts, test, false));
-			runs++;
-		}
-		found++;
-	}
-
-	if (test_name && !found)
-		printf("Test '%s' not found\n", test_name);
-	else
-		printf("Failures: %d\n", uts->fail_count);
-
-	/* Put everything back to normal so that sandbox works as expected */
-	gd_set_of_root(uts->of_root);
-	gd->dm_root = NULL;
-	ut_assertok(dm_init(CONFIG_IS_ENABLED(OF_LIVE)));
-	dm_scan_plat(false);
-	if (!CONFIG_IS_ENABLED(OF_PLATDATA))
-		dm_scan_fdt(false);
-
-	return uts->fail_count ? CMD_RET_FAILURE : 0;
-}
-
-int do_ut_dm(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
-{
-	const char *test_name = NULL;
-
-	if (argc > 1)
-		test_name = argv[1];
-
-	return dm_test_main(test_name);
-}
diff --git a/test/dm/test-uclass.c b/test/dm/test-uclass.c
index f1b7aaa727..067701734a 100644
--- a/test/dm/test-uclass.c
+++ b/test/dm/test-uclass.c
@@ -17,8 +17,6 @@
 #include <test/test.h>
 #include <test/ut.h>
 
-static struct unit_test_state *uts = &global_dm_test_state;
-
 int test_ping(struct udevice *dev, int pingval, int *pingret)
 {
 	const struct test_ops *ops = device_get_ops(dev);
@@ -31,6 +29,7 @@ int test_ping(struct udevice *dev, int pingval, int *pingret)
 
 static int test_post_bind(struct udevice *dev)
 {
+	struct unit_test_state *uts = test_get_state();
 	struct dm_test_perdev_uc_pdata *uc_pdata;
 
 	dm_testdrv_op_count[DM_TEST_OP_POST_BIND]++;
@@ -56,6 +55,7 @@ static int test_pre_unbind(struct udevice *dev)
 static int test_pre_probe(struct udevice *dev)
 {
 	struct dm_test_uclass_perdev_priv *priv = dev_get_uclass_priv(dev);
+	struct unit_test_state *uts = test_get_state();
 
 	dm_testdrv_op_count[DM_TEST_OP_PRE_PROBE]++;
 	ut_assert(priv);
@@ -66,18 +66,18 @@ static int test_pre_probe(struct udevice *dev)
 
 static int test_post_probe(struct udevice *dev)
 {
+	struct unit_test_state *uts = test_get_state();
 	struct udevice *prev = list_entry(dev->uclass_node.prev,
 					    struct udevice, uclass_node);
 
 	struct dm_test_uclass_perdev_priv *priv = dev_get_uclass_priv(dev);
 	struct uclass *uc = dev->uclass;
-	struct dm_test_state *dms = uts->priv;
 
 	dm_testdrv_op_count[DM_TEST_OP_POST_PROBE]++;
 	ut_assert(priv);
 	ut_assert(device_active(dev));
 	priv->base_add = 0;
-	if (dms->skip_post_probe)
+	if (uts->skip_post_probe)
 		return 0;
 	if (&prev->uclass_node != &uc->dev_head) {
 		struct dm_test_uclass_perdev_priv *prev_uc_priv
@@ -101,6 +101,8 @@ static int test_pre_remove(struct udevice *dev)
 
 static int test_init(struct uclass *uc)
 {
+	struct unit_test_state *uts = test_get_state();
+
 	dm_testdrv_op_count[DM_TEST_OP_INIT]++;
 	ut_assert(uclass_get_priv(uc));
 
diff --git a/test/env/cmd_ut_env.c b/test/env/cmd_ut_env.c
index a440b1bfb0..d65a32179c 100644
--- a/test/env/cmd_ut_env.c
+++ b/test/env/cmd_ut_env.c
@@ -12,8 +12,8 @@
 
 int do_ut_env(struct cmd_tbl *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 *tests = UNIT_TEST_SUITE_START(env_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(env_test);
 
 	return cmd_ut_category("environment", "env_test_",
 			       tests, n_ents, argc, argv);
diff --git a/test/image/Makefile b/test/image/Makefile
new file mode 100644
index 0000000000..c4039df707
--- /dev/null
+++ b/test/image/Makefile
@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: GPL-2.0+
+#
+# Copyright 2021 Google LLC
+
+obj-$(CONFIG_SPL_BUILD) += spl_load.o
diff --git a/test/image/spl_load.c b/test/image/spl_load.c
new file mode 100644
index 0000000000..851603ddd7
--- /dev/null
+++ b/test/image/spl_load.c
@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * Copyright 2021 Google LLC
+ * Written by Simon Glass <sjg@chromium.org>
+ */
+
+#include <common.h>
+#include <image.h>
+#include <mapmem.h>
+#include <os.h>
+#include <spl.h>
+#include <test/ut.h>
+
+/* Declare a new SPL test */
+#define SPL_TEST(_name, _flags)		UNIT_TEST(_name, _flags, spl_test)
+
+/* Context used for this test */
+struct text_ctx {
+	int fd;
+};
+
+static ulong read_fit_image(struct spl_load_info *load, ulong sector,
+			    ulong count, void *buf)
+{
+	struct text_ctx *text_ctx = load->priv;
+	off_t offset, ret;
+	ssize_t res;
+
+	offset = sector * load->bl_len;
+	ret = os_lseek(text_ctx->fd, offset, OS_SEEK_SET);
+	if (ret != offset) {
+		printf("Failed to seek to %zx, got %zx (errno=%d)\n", offset,
+		       ret, errno);
+		return 0;
+	}
+
+	res = os_read(text_ctx->fd, buf, count * load->bl_len);
+	if (res == -1) {
+		printf("Failed to read %lx bytes, got %ld (errno=%d)\n",
+		       count * load->bl_len, res, errno);
+		return 0;
+	}
+
+	return count;
+}
+
+int board_fit_config_name_match(const char *name)
+{
+	return 0;
+}
+
+struct image_header *spl_get_load_buffer(ssize_t offset, size_t size)
+{
+	return map_sysmem(0x100000, 0);
+}
+
+static int spl_test_load(struct unit_test_state *uts)
+{
+	struct spl_image_info image;
+	struct image_header *header;
+	struct text_ctx text_ctx;
+	struct spl_load_info load;
+	char fname[256];
+	int ret;
+	int fd;
+
+	memset(&load, '\0', sizeof(load));
+	load.bl_len = 512;
+	load.read = read_fit_image;
+
+	ret = os_find_u_boot(fname, sizeof(fname), true);
+	if (ret) {
+		printf("(%s not found, error %d)\n", fname, ret);
+		return ret;
+	}
+	load.filename = fname;
+
+	header = spl_get_load_buffer(-sizeof(*header), sizeof(*header));
+
+	fd = os_open(fname, OS_O_RDONLY);
+	ut_assert(fd >= 0);
+	ut_asserteq(512, os_read(fd, header, 512));
+	text_ctx.fd = fd;
+
+	load.priv = &text_ctx;
+
+	ut_assertok(spl_load_simple_fit(&image, &load, 0, header));
+
+	return 0;
+}
+SPL_TEST(spl_test_load, 0);
diff --git a/test/lib/cmd_ut_lib.c b/test/lib/cmd_ut_lib.c
index f5c7bf3d3b..f1ac015b2c 100644
--- a/test/lib/cmd_ut_lib.c
+++ b/test/lib/cmd_ut_lib.c
@@ -13,8 +13,8 @@
 
 int do_ut_lib(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test, lib_test);
-	const int n_ents = ll_entry_count(struct unit_test, lib_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(lib_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(lib_test);
 
 	return cmd_ut_category("lib", "lib_test_", tests, n_ents, argc, argv);
 }
diff --git a/test/log/Makefile b/test/log/Makefile
index 3f09deb644..a3dedace04 100644
--- a/test/log/Makefile
+++ b/test/log/Makefile
@@ -7,7 +7,7 @@ obj-$(CONFIG_CMD_LOG) += log_filter.o
 
 ifdef CONFIG_UT_LOG
 
-obj-y += test-main.o
+obj-y += log_ut.o
 
 ifdef CONFIG_SANDBOX
 obj-$(CONFIG_LOG_SYSLOG) += syslog_test.o
diff --git a/test/log/test-main.c b/test/log/log_ut.c
index c534add493..5aa3a18400 100644
--- a/test/log/test-main.c
+++ b/test/log/log_ut.c
@@ -13,8 +13,8 @@
 
 int do_ut_log(struct cmd_tbl *cmdtp, int flag, int argc, char * const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test, log_test);
-	const int n_ents = ll_entry_count(struct unit_test, log_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(log_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(log_test);
 
 	return cmd_ut_category("log", "log_test_",
 			       tests, n_ents, argc, argv);
diff --git a/test/optee/cmd_ut_optee.c b/test/optee/cmd_ut_optee.c
index 9fa4c91e0d..c3887ab11d 100644
--- a/test/optee/cmd_ut_optee.c
+++ b/test/optee/cmd_ut_optee.c
@@ -94,9 +94,8 @@ OPTEE_TEST(optee_fdt_protected_memory, 0);
 
 int do_ut_optee(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test,
-						 optee_test);
-	const int n_ents = ll_entry_count(struct unit_test, optee_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(optee_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(optee_test);
 	struct unit_test_state *uts;
 	void *fdt_optee = &__dtb_test_optee_optee_begin;
 	void *fdt_no_optee = &__dtb_test_optee_no_optee_begin;
diff --git a/test/overlay/cmd_ut_overlay.c b/test/overlay/cmd_ut_overlay.c
index c001fb183f..56a3df1713 100644
--- a/test/overlay/cmd_ut_overlay.c
+++ b/test/overlay/cmd_ut_overlay.c
@@ -213,9 +213,8 @@ OVERLAY_TEST(fdt_overlay_stacked, 0);
 
 int do_ut_overlay(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test,
-						 overlay_test);
-	const int n_ents = ll_entry_count(struct unit_test, overlay_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(overlay_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(overlay_test);
 	struct unit_test_state *uts;
 	void *fdt_base = &__dtb_test_fdt_base_begin;
 	void *fdt_overlay = &__dtb_test_fdt_overlay_begin;
diff --git a/test/py/conftest.py b/test/py/conftest.py
index 9bfd926345..1b909cde9d 100644
--- a/test/py/conftest.py
+++ b/test/py/conftest.py
@@ -226,7 +226,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_(.*)_test_2_\1_test_(.*)\s*$')
+re_ut_test_list = re.compile(r'_u_boot_list_2_ut_(.*)_test_2_\1_test_(.*)\s*$')
 def generate_ut_subtest(metafunc, fixture_name, sym_path):
     """Provide parametrization for a ut_subtest fixture.
 
diff --git a/test/py/tests/test_log.py b/test/py/tests/test_log.py
index f889120f2b..140dcb9aa2 100644
--- a/test/py/tests/test_log.py
+++ b/test/py/tests/test_log.py
@@ -45,5 +45,4 @@ def test_log_dropped(u_boot_console):
     cons = u_boot_console
     cons.restart_uboot()
     output = cons.get_spawn_output().replace('\r', '')
-    assert 'sandbox: starting...' in output
     assert (not 'debug: main' in output)
diff --git a/test/py/tests/test_ofplatdata.py b/test/py/tests/test_ofplatdata.py
index 92d09b7aa1..e9cce4daf4 100644
--- a/test/py/tests/test_ofplatdata.py
+++ b/test/py/tests/test_ofplatdata.py
@@ -4,7 +4,7 @@
 import pytest
 import u_boot_utils as util
 
-@pytest.mark.boardspec('sandbox')
+@pytest.mark.boardspec('sandbox_spl')
 @pytest.mark.buildconfigspec('spl_of_platdata')
 def test_spl_devicetree(u_boot_console):
     """Test content of spl device-tree"""
diff --git a/test/str_ut.c b/test/str_ut.c
index cd5045516d..359d7d4ea1 100644
--- a/test/str_ut.c
+++ b/test/str_ut.c
@@ -107,9 +107,8 @@ STR_TEST(str_simple_strtoul, 0);
 
 int do_ut_str(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test,
-						 str_test);
-	const int n_ents = ll_entry_count(struct unit_test, str_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(str_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(str_test);
 
 	return cmd_ut_category("str", "str_", tests, n_ents, argc, argv);
 }
diff --git a/test/test-main.c b/test/test-main.c
new file mode 100644
index 0000000000..e1b49e091a
--- /dev/null
+++ b/test/test-main.c
@@ -0,0 +1,426 @@
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * Copyright 2021 Google LLC
+ * Written by Simon Glass <sjg@chromium.org>
+ */
+
+#include <common.h>
+#include <console.h>
+#include <dm.h>
+#include <asm/state.h>
+#include <dm/root.h>
+#include <dm/test.h>
+#include <dm/uclass-internal.h>
+#include <test/test.h>
+#include <test/ut.h>
+
+DECLARE_GLOBAL_DATA_PTR;
+
+/* This is valid when a test is running, NULL otherwise */
+static struct unit_test_state *cur_test_state;
+
+struct unit_test_state *test_get_state(void)
+{
+	return cur_test_state;
+}
+
+void test_set_state(struct unit_test_state *uts)
+{
+	cur_test_state = uts;
+}
+
+/**
+ * dm_test_pre_run() - Get ready to run a driver model test
+ *
+ * This clears out the driver model data structures. For sandbox it resets the
+ * state structure
+ *
+ * @uts: Test state
+ */
+static int dm_test_pre_run(struct unit_test_state *uts)
+{
+	bool of_live = uts->of_live;
+
+	uts->root = NULL;
+	uts->testdev = NULL;
+	uts->force_fail_alloc = false;
+	uts->skip_post_probe = false;
+	gd->dm_root = NULL;
+	if (!CONFIG_IS_ENABLED(OF_PLATDATA))
+		memset(dm_testdrv_op_count, '\0', sizeof(dm_testdrv_op_count));
+	state_reset_for_test(state_get_current());
+
+	/* Determine whether to make the live tree available */
+	gd_set_of_root(of_live ? uts->of_root : NULL);
+	ut_assertok(dm_init(of_live));
+	uts->root = dm_root();
+
+	return 0;
+}
+
+static int dm_test_post_run(struct unit_test_state *uts)
+{
+	int id;
+
+	for (id = 0; id < UCLASS_COUNT; id++) {
+		struct uclass *uc;
+
+		/*
+		 * If the uclass doesn't exist we don't want to create it. So
+		 * check that here before we call uclass_find_device().
+		 */
+		uc = uclass_find(id);
+		if (!uc)
+			continue;
+		ut_assertok(uclass_destroy(uc));
+	}
+
+	return 0;
+}
+
+/* Ensure all the test devices are probed */
+static int do_autoprobe(struct unit_test_state *uts)
+{
+	struct udevice *dev;
+	int ret;
+
+	/* Scanning the uclass is enough to probe all the devices */
+	for (ret = uclass_first_device(UCLASS_TEST, &dev);
+	     dev;
+	     ret = uclass_next_device(&dev))
+		;
+
+	return ret;
+}
+
+/*
+ * ut_test_run_on_flattree() - Check if we should run a test with flat DT
+ *
+ * This skips long/slow tests where there is not much value in running a flat
+ * DT test in addition to a live DT test.
+ *
+ * @return true to run the given test on the flat device tree
+ */
+static bool ut_test_run_on_flattree(struct unit_test *test)
+{
+	const char *fname = strrchr(test->file, '/') + 1;
+
+	if (!(test->flags & UT_TESTF_DM))
+		return false;
+
+	return !strstr(fname, "video") || strstr(test->name, "video_base");
+}
+
+/**
+ * test_matches() - Check if a test should be run
+ *
+ * This checks if the a test should be run. In the normal case of running all
+ * tests, @select_name is NULL.
+ *
+ * @prefix: String prefix for the tests. Any tests that have this prefix will be
+ *	printed without the prefix, so that it is easier to see the unique part
+ *	of the test name. If NULL, any suite name (xxx_test) is considered to be
+ *	a prefix.
+ * @test_name: Name of current test
+ * @select_name: Name of test to run (or NULL for all)
+ * @return true to run this test, false to skip it
+ */
+static bool test_matches(const char *prefix, const char *test_name,
+			 const char *select_name)
+{
+	if (!select_name)
+		return true;
+
+	if (!strcmp(test_name, select_name))
+		return true;
+
+	if (!prefix) {
+		const char *p = strstr(test_name, "_test_");
+
+		/* convert xxx_test_yyy to yyy, i.e. remove the suite name */
+		if (p)
+			test_name = p + 6;
+	} else {
+		/* All tests have this prefix */
+		if (!strncmp(test_name, prefix, strlen(prefix)))
+			test_name += strlen(prefix);
+	}
+
+	if (!strcmp(test_name, select_name))
+		return true;
+
+	return false;
+}
+
+/**
+ * ut_list_has_dm_tests() - Check if a list of tests has driver model ones
+ *
+ * @tests: List of tests to run
+ * @count: Number of tests to ru
+ * @return true if any of the tests have the UT_TESTF_DM flag
+ */
+static bool ut_list_has_dm_tests(struct unit_test *tests, int count)
+{
+	struct unit_test *test;
+
+	for (test = tests; test < tests + count; test++) {
+		if (test->flags & UT_TESTF_DM)
+			return true;
+	}
+
+	return false;
+}
+
+/**
+ * dm_test_restore() Put things back to normal so sandbox works as expected
+ *
+ * @of_root: Value to set for of_root
+ * @return 0 if OK, -ve on error
+ */
+static int dm_test_restore(struct device_node *of_root)
+{
+	int ret;
+
+	gd_set_of_root(of_root);
+	gd->dm_root = NULL;
+	ret = dm_init(CONFIG_IS_ENABLED(OF_LIVE));
+	if (ret)
+		return ret;
+	dm_scan_plat(false);
+	if (!CONFIG_IS_ENABLED(OF_PLATDATA))
+		dm_scan_fdt(false);
+
+	return 0;
+}
+
+/**
+ * test_pre_run() - Handle any preparation needed to run a test
+ *
+ * @uts: Test state
+ * @test: Test to prepare for
+ * @return 0 if OK, -EAGAIN to skip this test since some required feature is not
+ *	available, other -ve on error (meaning that testing cannot likely
+ *	continue)
+ */
+static int test_pre_run(struct unit_test_state *uts, struct unit_test *test)
+{
+	if (test->flags & UT_TESTF_DM)
+		ut_assertok(dm_test_pre_run(uts));
+
+	ut_set_skip_delays(uts, false);
+
+	uts->start = mallinfo();
+
+	if (test->flags & UT_TESTF_SCAN_PDATA)
+		ut_assertok(dm_scan_plat(false));
+
+	if (test->flags & UT_TESTF_PROBE_TEST)
+		ut_assertok(do_autoprobe(uts));
+
+	if (!CONFIG_IS_ENABLED(OF_PLATDATA) &&
+	    (test->flags & UT_TESTF_SCAN_FDT))
+		ut_assertok(dm_extended_scan(false));
+
+	if (test->flags & UT_TESTF_CONSOLE_REC) {
+		int ret = console_record_reset_enable();
+
+		if (ret) {
+			printf("Skipping: Console recording disabled\n");
+			return -EAGAIN;
+		}
+	}
+	ut_silence_console(uts);
+
+	return 0;
+}
+
+/**
+ * test_post_run() - Handle cleaning up after a test
+ *
+ * @uts: Test state
+ * @test: Test to clean up after
+ * @return 0 if OK, -ve on error (meaning that testing cannot likely continue)
+ */
+static int test_post_run(struct unit_test_state *uts, struct unit_test *test)
+{
+	ut_unsilence_console(uts);
+	if (test->flags & UT_TESTF_DM)
+		ut_assertok(dm_test_post_run(uts));
+
+	return 0;
+}
+
+/**
+ * ut_run_test() - Run a single test
+ *
+ * This runs the test, handling any preparation and clean-up needed. It prints
+ * the name of each test before running it.
+ *
+ * @uts: Test state to update. The caller should ensure that this is zeroed for
+ *	the first call to this function. On exit, @uts->fail_count is
+ *	incremented by the number of failures (0, one hopes)
+ * @test_name: Test to run
+ * @name: Name of test, possibly skipping a prefix that should not be displayed
+ * @return 0 if all tests passed, -EAGAIN if the test should be skipped, -1 if
+ *	any failed
+ */
+static int ut_run_test(struct unit_test_state *uts, struct unit_test *test,
+		       const char *test_name)
+{
+	const char *fname = strrchr(test->file, '/') + 1;
+	const char *note = "";
+	int ret;
+
+	if ((test->flags & UT_TESTF_DM) && !uts->of_live)
+		note = " (flat tree)";
+	printf("Test: %s: %s%s\n", test_name, fname, note);
+
+	/* Allow access to test state from drivers */
+	test_set_state(uts);
+
+	ret = test_pre_run(uts, test);
+	if (ret == -EAGAIN)
+		return -EAGAIN;
+	if (ret)
+		return ret;
+
+	test->func(uts);
+
+	ret = test_post_run(uts, test);
+	if (ret)
+		return ret;
+
+	test_set_state( NULL);
+
+	return 0;
+}
+
+/**
+ * ut_run_test_live_flat() - Run a test with both live and flat tree
+ *
+ * This calls ut_run_test() with livetree enabled, which is the standard setup
+ * for runnig tests. Then, for driver model test, it calls it again with
+ * livetree disabled. This allows checking of flattree being used when OF_LIVE
+ * is enabled, as is the case in U-Boot proper before relocation, as well as in
+ * SPL.
+ *
+ * @uts: Test state to update. The caller should ensure that this is zeroed for
+ *	the first call to this function. On exit, @uts->fail_count is
+ *	incremented by the number of failures (0, one hopes)
+ * @test: Test to run
+ * @name: Name of test, possibly skipping a prefix that should not be displayed
+ * @return 0 if all tests passed, -EAGAIN if the test should be skipped, -1 if
+ *	any failed
+ */
+static int ut_run_test_live_flat(struct unit_test_state *uts,
+				 struct unit_test *test, const char *name)
+{
+	int runs;
+
+	/* Run with the live tree if possible */
+	runs = 0;
+	if (CONFIG_IS_ENABLED(OF_LIVE)) {
+		if (!(test->flags & UT_TESTF_FLAT_TREE)) {
+			uts->of_live = true;
+			ut_assertok(ut_run_test(uts, test, test->name));
+			runs++;
+		}
+	}
+
+	/*
+	 * Run with the flat tree if we couldn't run it with live tree,
+	 * or it is a core test.
+	 */
+	if (!(test->flags & UT_TESTF_LIVE_TREE) &&
+	    (!runs || ut_test_run_on_flattree(test))) {
+		uts->of_live = false;
+		ut_assertok(ut_run_test(uts, test, test->name));
+		runs++;
+	}
+
+	return 0;
+}
+
+/**
+ * ut_run_tests() - Run a set of tests
+ *
+ * This runs the tests, handling any preparation and clean-up needed. It prints
+ * the name of each test before running it.
+ *
+ * @uts: Test state to update. The caller should ensure that this is zeroed for
+ *	the first call to this function. On exit, @uts->fail_count is
+ *	incremented by the number of failures (0, one hopes)
+ * @prefix: String prefix for the tests. Any tests that have this prefix will be
+ *	printed without the prefix, so that it is easier to see the unique part
+ *	of the test name. If NULL, no prefix processing is done
+ * @tests: List of tests to run
+ * @count: Number of tests to run
+ * @select_name: Name of a single test to run (from the list provided). If NULL
+ *	then all tests are run
+ * @return 0 if all tests passed, -ENOENT if test @select_name was not found,
+ *	-EBADF if any failed
+ */
+static int ut_run_tests(struct unit_test_state *uts, const char *prefix,
+			struct unit_test *tests, int count,
+			const char *select_name)
+{
+	struct unit_test *test;
+	int found = 0;
+
+	for (test = tests; test < tests + count; test++) {
+		const char *test_name = test->name;
+		int ret;
+
+		if (!test_matches(prefix, test_name, select_name))
+			continue;
+		ret = ut_run_test_live_flat(uts, test, select_name);
+		found++;
+		if (ret == -EAGAIN)
+			continue;
+		if (ret)
+			return ret;
+	}
+	if (select_name && !found)
+		return -ENOENT;
+
+	return uts->fail_count ? -EBADF : 0;
+}
+
+int ut_run_list(const char *category, const char *prefix,
+		struct unit_test *tests, int count, const char *select_name)
+{
+	struct unit_test_state uts = { .fail_count = 0 };
+	bool has_dm_tests = false;
+	int ret;
+
+	if (!CONFIG_IS_ENABLED(OF_PLATDATA) &&
+	    ut_list_has_dm_tests(tests, count)) {
+		has_dm_tests = true;
+		/*
+		 * If we have no device tree, or it only has a root node, then
+		 * these * tests clearly aren't going to work...
+		 */
+		if (!gd->fdt_blob || fdt_next_node(gd->fdt_blob, 0, NULL) < 0) {
+			puts("Please run with test device tree:\n"
+			     "    ./u-boot -d arch/sandbox/dts/test.dtb\n");
+			return CMD_RET_FAILURE;
+		}
+	}
+
+	if (!select_name)
+		printf("Running %d %s tests\n", count, category);
+
+	uts.of_root = gd_of_root();
+	ret = ut_run_tests(&uts, prefix, tests, count, select_name);
+
+	if (ret == -ENOENT)
+		printf("Test '%s' not found\n", select_name);
+	else
+		printf("Failures: %d\n", uts.fail_count);
+
+	/* Best efforts only...ignore errors */
+	if (has_dm_tests)
+		dm_test_restore(uts.of_root);
+
+	return ret;
+}
diff --git a/test/unicode_ut.c b/test/unicode_ut.c
index 6130ef0b54..7b9c020eff 100644
--- a/test/unicode_ut.c
+++ b/test/unicode_ut.c
@@ -615,8 +615,8 @@ UNICODE_TEST(unicode_test_efi_create_indexed_name);
 
 int do_ut_unicode(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
 {
-	struct unit_test *tests = ll_entry_start(struct unit_test, unicode_test);
-	const int n_ents = ll_entry_count(struct unit_test, unicode_test);
+	struct unit_test *tests = UNIT_TEST_SUITE_START(unicode_test);
+	const int n_ents = UNIT_TEST_SUITE_COUNT(unicode_test);
 
 	return cmd_ut_category("Unicode", "unicode_test_",
 			       tests, n_ents, argc, argv);
diff --git a/test/ut.c b/test/ut.c
index 7328338731..ea0af153e4 100644
--- a/test/ut.c
+++ b/test/ut.c
@@ -133,3 +133,10 @@ void ut_unsilence_console(struct unit_test_state *uts)
 {
 	gd->flags &= ~(GD_FLG_SILENT | GD_FLG_RECORD);
 }
+
+void ut_set_skip_delays(struct unit_test_state *uts, bool skip_delays)
+{
+#ifdef CONFIG_SANDBOX
+	state_set_skip_delays(skip_delays);
+#endif
+}