test: Generalize the unit test framework

Separate the ability to define tests and assert status of test functions
from the dm tests so they can be used more consistently throughout all
tests.

Signed-off-by: Joe Hershberger <joe.hershberger@ni.com>
Reviewed-by: Simon Glass <sjg@chromium.org>

24 files changed, 234 insertions, 191 deletions

diff --git a/include/dm/test.h b/include/dm/test.h
index f03fbcb..98f2b9e 100644
--- a/include/dm/test.h
+++ b/include/dm/test.h
@@ -8,7 +8,7 @@
 #define __DM_TEST_H
 
 #include <dm.h>
-#include <malloc.h>
+#include <test/test.h>
 
 /**
  * struct dm_test_cdata - configuration data for test instance
@@ -124,7 +124,7 @@ struct dm_test_perdev_uc_pdata {
  */
 extern int dm_testdrv_op_count[DM_TEST_OP_COUNT];
 
-extern struct dm_test_state global_test_state;
+extern struct unit_test_state global_dm_test_state;
 
 /*
  * struct dm_test_state - Entire state of dm test system
@@ -133,7 +133,6 @@ extern struct dm_test_state global_test_state;
  *
  * @root: Root device
  * @testdev: Test device
- * @fail_count: Number of tests that failed
  * @force_fail_alloc: Force all memory allocs to fail
  * @skip_post_probe: Skip uclass post-probe processing
  * @removed: Used to keep track of a device that was removed
@@ -141,11 +140,9 @@ extern struct dm_test_state global_test_state;
 struct dm_test_state {
 	struct udevice *root;
 	struct udevice *testdev;
-	int fail_count;
 	int force_fail_alloc;
 	int skip_post_probe;
 	struct udevice *removed;
-	struct mallinfo start;
 };
 
 /* Test flags for each test */
@@ -155,26 +152,8 @@ enum {
 	DM_TESTF_SCAN_FDT	= 1 << 2,	/* scan device tree */
 };
 
-/**
- * struct dm_test - Information about a driver model test
- *
- * @name: Name of test
- * @func: Function to call to perform test
- * @flags: Flags indicated pre-conditions for test
- */
-struct dm_test {
-	const char *name;
-	int (*func)(struct dm_test_state *dms);
-	int flags;
-};
-
 /* Declare a new driver model test */
-#define DM_TEST(_name, _flags)						\
-	ll_entry_declare(struct dm_test, _name, dm_test) = {		\
-		.name = #_name,						\
-		.flags = _flags,					\
-		.func = _name,						\
-	}
+#define DM_TEST(_name, _flags)	UNIT_TEST(_name, _flags, dm_test)
 
 /* Declare ping methods for the drivers */
 int test_ping(struct udevice *dev, int pingval, int *pingret);
@@ -191,7 +170,7 @@ int testfdt_ping(struct udevice *dev, int pingval, int *pingret);
  * @priv: Pointer to private test information
  * @return 0 if OK, -ve on error
  */
-int dm_check_operations(struct dm_test_state *dms, struct udevice *dev,
+int dm_check_operations(struct unit_test_state *uts, struct udevice *dev,
 			uint32_t base, struct dm_test_priv *priv);
 
 /**
@@ -201,7 +180,7 @@ int dm_check_operations(struct dm_test_state *dms, struct udevice *dev,
  * @num_devices: Number of test devices to check
  * @return 0 if OK, -ve on error
  */
-int dm_check_devices(struct dm_test_state *dms, int num_devices);
+int dm_check_devices(struct unit_test_state *uts, int num_devices);
 
 /**
  * dm_leak_check_start() - Prepare to check for a memory leak
@@ -211,7 +190,7 @@ int dm_check_devices(struct dm_test_state *dms, int num_devices);
  *
  * @dms: Overall test state
  */
-void dm_leak_check_start(struct dm_test_state *dms);
+void dm_leak_check_start(struct unit_test_state *uts);
 
 /**
  * dm_leak_check_end() - Check that no memory has leaked
@@ -221,7 +200,7 @@ void dm_leak_check_start(struct dm_test_state *dms);
  * it sees a different amount of total memory allocated than before.
  *
  * @dms: Overall test state
- */int dm_leak_check_end(struct dm_test_state *dms);
+ */int dm_leak_check_end(struct unit_test_state *uts);
 
 
 /**
diff --git a/include/test/test.h b/include/test/test.h
new file mode 100644
index 0000000..b7e1ae2
--- /dev/null
+++ b/include/test/test.h
@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) 2013 Google, Inc.
+ *
+ * SPDX-License-Identifier:	GPL-2.0+
+ */
+
+#ifndef __TEST_TEST_H
+#define __TEST_TEST_H
+
+#include <malloc.h>
+
+/*
+ * struct unit_test_state - Entire state of test system
+ *
+ * @fail_count: Number of tests that failed
+ * @start: Store the starting mallinfo when doing leak test
+ * @priv: A pointer to some other info some suites want to track
+ */
+struct unit_test_state {
+	int fail_count;
+	struct mallinfo start;
+	void *priv;
+};
+
+/**
+ * struct unit_test - Information about a unit test
+ *
+ * @name: Name of test
+ * @func: Function to call to perform test
+ * @flags: Flags indicated pre-conditions for test
+ */
+struct unit_test {
+	const char *name;
+	int (*func)(struct unit_test_state *state);
+	int flags;
+};
+
+/* Declare a new unit test */
+#define UNIT_TEST(_name, _flags, _suite)				\
+	ll_entry_declare(struct unit_test, _name, _suite) = {		\
+		.name = #_name,						\
+		.flags = _flags,					\
+		.func = _name,						\
+	}
+
+
+#endif /* __TEST_TEST_H */
diff --git a/include/dm/ut.h b/include/test/ut.h
index ec61465..275f27f 100644
--- a/include/dm/ut.h
+++ b/include/test/ut.h
@@ -1,39 +1,39 @@
 /*
- * Simple unit test library for driver model
+ * Simple unit test library
  *
  * Copyright (c) 2013 Google, Inc
  *
  * SPDX-License-Identifier:	GPL-2.0+
  */
 
-#ifndef __DM_UT_H
-#define __DM_UT_H
+#ifndef __TEST_UT_H
+#define __TEST_UT_H
 
-struct dm_test_state;
+struct unit_test_state;
 
 /**
  * ut_fail() - Record failure of a unit test
  *
- * @dms: Test state
+ * @uts: Test state
  * @fname: Filename where the error occured
  * @line: Line number where the error occured
  * @func: Function name where the error occured
  * @cond: The condition that failed
  */
-void ut_fail(struct dm_test_state *dms, const char *fname, int line,
+void ut_fail(struct unit_test_state *uts, const char *fname, int line,
 	     const char *func, const char *cond);
 
 /**
  * ut_failf() - Record failure of a unit test
  *
- * @dms: Test state
+ * @uts: Test state
  * @fname: Filename where the error occured
  * @line: Line number where the error occured
  * @func: Function name where the error occured
  * @cond: The condition that failed
  * @fmt: printf() format string for the error, followed by args
  */
-void ut_failf(struct dm_test_state *dms, const char *fname, int line,
+void ut_failf(struct unit_test_state *uts, const char *fname, int line,
 	      const char *func, const char *cond, const char *fmt, ...)
 			__attribute__ ((format (__printf__, 6, 7)));
 
@@ -41,14 +41,14 @@ void ut_failf(struct dm_test_state *dms, const char *fname, int line,
 /* Assert that a condition is non-zero */
 #define ut_assert(cond)							\
 	if (!(cond)) {							\
-		ut_fail(dms, __FILE__, __LINE__, __func__, #cond);	\
+		ut_fail(uts, __FILE__, __LINE__, __func__, #cond);	\
 		return -1;						\
 	}
 
 /* Assert that a condition is non-zero, with printf() string */
 #define ut_assertf(cond, fmt, args...)					\
 	if (!(cond)) {							\
-		ut_failf(dms, __FILE__, __LINE__, __func__, #cond,	\
+		ut_failf(uts, __FILE__, __LINE__, __func__, #cond,	\
 			 fmt, ##args);					\
 		return -1;						\
 	}
@@ -58,7 +58,7 @@ void ut_failf(struct dm_test_state *dms, const char *fname, int line,
 	unsigned int val1 = (expr1), val2 = (expr2);			\
 									\
 	if (val1 != val2) {						\
-		ut_failf(dms, __FILE__, __LINE__, __func__,		\
+		ut_failf(uts, __FILE__, __LINE__, __func__,		\
 			 #expr1 " == " #expr2,				\
 			 "Expected %d, got %d", val1, val2);		\
 		return -1;						\
@@ -70,7 +70,7 @@ void ut_failf(struct dm_test_state *dms, const char *fname, int line,
 	const char *val1 = (expr1), *val2 = (expr2);			\
 									\
 	if (strcmp(val1, val2)) {					\
-		ut_failf(dms, __FILE__, __LINE__, __func__,		\
+		ut_failf(uts, __FILE__, __LINE__, __func__,		\
 			 #expr1 " = " #expr2,				\
 			 "Expected \"%s\", got \"%s\"", val1, val2);	\
 		return -1;						\
@@ -82,7 +82,7 @@ void ut_failf(struct dm_test_state *dms, const char *fname, int line,
 	const void *val1 = (expr1), *val2 = (expr2);			\
 									\
 	if (val1 != val2) {						\
-		ut_failf(dms, __FILE__, __LINE__, __func__,		\
+		ut_failf(uts, __FILE__, __LINE__, __func__,		\
 			 #expr1 " = " #expr2,				\
 			 "Expected %p, got %p", val1, val2);		\
 		return -1;						\
@@ -94,7 +94,7 @@ void ut_failf(struct dm_test_state *dms, const char *fname, int line,
 	const void *val = (expr);					\
 									\
 	if (val == NULL) {						\
-		ut_failf(dms, __FILE__, __LINE__, __func__,		\
+		ut_failf(uts, __FILE__, __LINE__, __func__,		\
 			 #expr " = NULL",				\
 			 "Expected non-null, got NULL");		\
 		return -1;						\
diff --git a/test/Kconfig b/test/Kconfig
index 3270c84..6f918ed 100644
--- a/test/Kconfig
+++ b/test/Kconfig
@@ -1,3 +1,6 @@
+config UNIT_TEST
+	bool
+
 config CMD_UT_TIME
 	bool "Unit tests for time functions"
 	help
diff --git a/test/Makefile b/test/Makefile
index 08330e0..422f08f 100644
--- a/test/Makefile
+++ b/test/Makefile
@@ -4,6 +4,7 @@
 # SPDX-License-Identifier:	GPL-2.0+
 #
 
+obj-$(CONFIG_UNIT_TEST) += ut.o
 obj-$(CONFIG_SANDBOX) += command_ut.o
 obj-$(CONFIG_SANDBOX) += compression.o
 obj-$(CONFIG_CMD_UT_TIME) += time_ut.o
diff --git a/test/dm/Kconfig b/test/dm/Kconfig
index a9d0298..3ca154f 100644
--- a/test/dm/Kconfig
+++ b/test/dm/Kconfig
@@ -1,6 +1,7 @@
 config DM_TEST
 	bool "Enable driver model test command"
 	depends on SANDBOX && CMD_DM
+	select UNIT_TEST
 	help
 	  This enables the 'dm test' command which runs a series of unit
 	  tests on the driver model code. Each subsystem (uclass) is tested.
diff --git a/test/dm/Makefile b/test/dm/Makefile
index c7087bb..07e782e 100644
--- a/test/dm/Makefile
+++ b/test/dm/Makefile
@@ -10,12 +10,10 @@ obj-$(CONFIG_DM_TEST) += test-driver.o
 obj-$(CONFIG_DM_TEST) += test-fdt.o
 obj-$(CONFIG_DM_TEST) += test-main.o
 obj-$(CONFIG_DM_TEST) += test-uclass.o
-obj-$(CONFIG_DM_TEST) += ut.o
 
 # Tests for particular subsystems - when enabling driver model for a new
 # subsystem you must add sandbox tests here.
 obj-$(CONFIG_DM_TEST) += core.o
-obj-$(CONFIG_DM_TEST) += ut.o
 ifneq ($(CONFIG_SANDBOX),)
 obj-$(CONFIG_DM_ETH) += eth.o
 obj-$(CONFIG_DM_GPIO) += gpio.o
diff --git a/test/dm/bus.c b/test/dm/bus.c
index 116a52d..a215905 100644
--- a/test/dm/bus.c
+++ b/test/dm/bus.c
@@ -10,8 +10,8 @@
 #include <dm/root.h>
 #include <dm/test.h>
 #include <dm/uclass-internal.h>
-#include <dm/ut.h>
 #include <dm/util.h>
+#include <test/ut.h>
 
 DECLARE_GLOBAL_DATA_PTR;
 
@@ -104,7 +104,7 @@ UCLASS_DRIVER(testbus) = {
 };
 
 /* Test that we can probe for children */
-static int dm_test_bus_children(struct dm_test_state *dms)
+static int dm_test_bus_children(struct unit_test_state *uts)
 {
 	int num_devices = 6;
 	struct udevice *bus;
@@ -120,14 +120,14 @@ static int dm_test_bus_children(struct dm_test_state *dms)
 	ut_assertok(uclass_get(UCLASS_TEST_FDT, &uc));
 	ut_asserteq(num_devices, list_count_items(&uc->dev_head));
 
-	ut_assert(!dm_check_devices(dms, num_devices));
+	ut_assert(!dm_check_devices(uts, num_devices));
 
 	return 0;
 }
 DM_TEST(dm_test_bus_children, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test our functions for accessing children */
-static int dm_test_bus_children_funcs(struct dm_test_state *dms)
+static int dm_test_bus_children_funcs(struct unit_test_state *uts)
 {
 	const void *blob = gd->fdt_blob;
 	struct udevice *bus, *dev;
@@ -173,7 +173,7 @@ static int dm_test_bus_children_funcs(struct dm_test_state *dms)
 DM_TEST(dm_test_bus_children_funcs, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that we can iterate through children */
-static int dm_test_bus_children_iterators(struct dm_test_state *dms)
+static int dm_test_bus_children_iterators(struct unit_test_state *uts)
 {
 	struct udevice *bus, *dev, *child;
 
@@ -204,7 +204,7 @@ DM_TEST(dm_test_bus_children_iterators,
 	DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that the bus can store data about each child */
-static int test_bus_parent_data(struct dm_test_state *dms)
+static int test_bus_parent_data(struct unit_test_state *uts)
 {
 	struct dm_test_parent_data *parent_data;
 	struct udevice *bus, *dev;
@@ -264,14 +264,14 @@ static int test_bus_parent_data(struct dm_test_state *dms)
 	return 0;
 }
 /* Test that the bus can store data about each child */
-static int dm_test_bus_parent_data(struct dm_test_state *dms)
+static int dm_test_bus_parent_data(struct unit_test_state *uts)
 {
-	return test_bus_parent_data(dms);
+	return test_bus_parent_data(uts);
 }
 DM_TEST(dm_test_bus_parent_data, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* As above but the size is controlled by the uclass */
-static int dm_test_bus_parent_data_uclass(struct dm_test_state *dms)
+static int dm_test_bus_parent_data_uclass(struct unit_test_state *uts)
 {
 	struct driver *drv;
 	struct udevice *bus;
@@ -284,7 +284,7 @@ static int dm_test_bus_parent_data_uclass(struct dm_test_state *dms)
 	size = drv->per_child_auto_alloc_size;
 	bus->uclass->uc_drv->per_child_auto_alloc_size = size;
 	drv->per_child_auto_alloc_size = 0;
-	ret = test_bus_parent_data(dms);
+	ret = test_bus_parent_data(uts);
 	if (ret)
 		return ret;
 	bus->uclass->uc_drv->per_child_auto_alloc_size = 0;
@@ -296,9 +296,10 @@ DM_TEST(dm_test_bus_parent_data_uclass,
 	DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that the bus ops are called when a child is probed/removed */
-static int dm_test_bus_parent_ops(struct dm_test_state *dms)
+static int dm_test_bus_parent_ops(struct unit_test_state *uts)
 {
 	struct dm_test_parent_data *parent_data;
+	struct dm_test_state *dms = uts->priv;
 	struct udevice *bus, *dev;
 	struct uclass *uc;
 
@@ -333,7 +334,7 @@ static int dm_test_bus_parent_ops(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_bus_parent_ops, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
-static int test_bus_parent_platdata(struct dm_test_state *dms)
+static int test_bus_parent_platdata(struct unit_test_state *uts)
 {
 	struct dm_test_parent_platdata *plat;
 	struct udevice *bus, *dev;
@@ -406,14 +407,14 @@ static int test_bus_parent_platdata(struct dm_test_state *dms)
 }
 
 /* Test that the bus can store platform data about each child */
-static int dm_test_bus_parent_platdata(struct dm_test_state *dms)
+static int dm_test_bus_parent_platdata(struct unit_test_state *uts)
 {
-	return test_bus_parent_platdata(dms);
+	return test_bus_parent_platdata(uts);
 }
 DM_TEST(dm_test_bus_parent_platdata, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* As above but the size is controlled by the uclass */
-static int dm_test_bus_parent_platdata_uclass(struct dm_test_state *dms)
+static int dm_test_bus_parent_platdata_uclass(struct unit_test_state *uts)
 {
 	struct udevice *bus;
 	struct driver *drv;
@@ -426,7 +427,7 @@ static int dm_test_bus_parent_platdata_uclass(struct dm_test_state *dms)
 	size = drv->per_child_platdata_auto_alloc_size;
 	bus->uclass->uc_drv->per_child_platdata_auto_alloc_size = size;
 	drv->per_child_platdata_auto_alloc_size = 0;
-	ret = test_bus_parent_platdata(dms);
+	ret = test_bus_parent_platdata(uts);
 	if (ret)
 		return ret;
 	bus->uclass->uc_drv->per_child_platdata_auto_alloc_size = 0;
@@ -438,7 +439,7 @@ DM_TEST(dm_test_bus_parent_platdata_uclass,
 	DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that the child post_bind method is called */
-static int dm_test_bus_child_post_bind(struct dm_test_state *dms)
+static int dm_test_bus_child_post_bind(struct unit_test_state *uts)
 {
 	struct dm_test_parent_platdata *plat;
 	struct udevice *bus, *dev;
@@ -461,7 +462,7 @@ static int dm_test_bus_child_post_bind(struct dm_test_state *dms)
 DM_TEST(dm_test_bus_child_post_bind, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that the child post_bind method is called */
-static int dm_test_bus_child_post_bind_uclass(struct dm_test_state *dms)
+static int dm_test_bus_child_post_bind_uclass(struct unit_test_state *uts)
 {
 	struct dm_test_parent_platdata *plat;
 	struct udevice *bus, *dev;
@@ -488,7 +489,7 @@ DM_TEST(dm_test_bus_child_post_bind_uclass,
  * Test that the bus' uclass' child_pre_probe() is called before the
  * device's probe() method
  */
-static int dm_test_bus_child_pre_probe_uclass(struct dm_test_state *dms)
+static int dm_test_bus_child_pre_probe_uclass(struct unit_test_state *uts)
 {
 	struct udevice *bus, *dev;
 	int child_count;
diff --git a/test/dm/core.c b/test/dm/core.c
index 91be1e5..976a706 100644
--- a/test/dm/core.c
+++ b/test/dm/core.c
@@ -13,10 +13,10 @@
 #include <malloc.h>
 #include <dm/device-internal.h>
 #include <dm/root.h>
-#include <dm/ut.h>
 #include <dm/util.h>
 #include <dm/test.h>
 #include <dm/uclass-internal.h>
+#include <test/ut.h>
 
 DECLARE_GLOBAL_DATA_PTR;
 
@@ -67,14 +67,14 @@ static struct driver_info driver_info_pre_reloc = {
 	.platdata = &test_pdata_manual,
 };
 
-void dm_leak_check_start(struct dm_test_state *dms)
+void dm_leak_check_start(struct unit_test_state *uts)
 {
-	dms->start = mallinfo();
-	if (!dms->start.uordblks)
+	uts->start = mallinfo();
+	if (!uts->start.uordblks)
 		puts("Warning: Please add '#define DEBUG' to the top of common/dlmalloc.c\n");
 }
 
-int dm_leak_check_end(struct dm_test_state *dms)
+int dm_leak_check_end(struct unit_test_state *uts)
 {
 	struct mallinfo end;
 	int id;
@@ -90,14 +90,15 @@ int dm_leak_check_end(struct dm_test_state *dms)
 	}
 
 	end = mallinfo();
-	ut_asserteq(dms->start.uordblks, end.uordblks);
+	ut_asserteq(uts->start.uordblks, end.uordblks);
 
 	return 0;
 }
 
 /* Test that binding with platdata occurs correctly */
-static int dm_test_autobind(struct dm_test_state *dms)
+static int dm_test_autobind(struct unit_test_state *uts)
 {
+	struct dm_test_state *dms = uts->priv;
 	struct udevice *dev;
 
 	/*
@@ -130,7 +131,7 @@ static int dm_test_autobind(struct dm_test_state *dms)
 DM_TEST(dm_test_autobind, 0);
 
 /* Test that binding with uclass platdata allocation occurs correctly */
-static int dm_test_autobind_uclass_pdata_alloc(struct dm_test_state *dms)
+static int dm_test_autobind_uclass_pdata_alloc(struct unit_test_state *uts)
 {
 	struct dm_test_perdev_uc_pdata *uc_pdata;
 	struct udevice *dev;
@@ -159,7 +160,7 @@ static int dm_test_autobind_uclass_pdata_alloc(struct dm_test_state *dms)
 DM_TEST(dm_test_autobind_uclass_pdata_alloc, DM_TESTF_SCAN_PDATA);
 
 /* Test that binding with uclass platdata setting occurs correctly */
-static int dm_test_autobind_uclass_pdata_valid(struct dm_test_state *dms)
+static int dm_test_autobind_uclass_pdata_valid(struct unit_test_state *uts)
 {
 	struct dm_test_perdev_uc_pdata *uc_pdata;
 	struct udevice *dev;
@@ -185,8 +186,9 @@ static int dm_test_autobind_uclass_pdata_valid(struct dm_test_state *dms)
 DM_TEST(dm_test_autobind_uclass_pdata_valid, DM_TESTF_SCAN_PDATA);
 
 /* Test that autoprobe finds all the expected devices */
-static int dm_test_autoprobe(struct dm_test_state *dms)
+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;
@@ -252,7 +254,7 @@ static int dm_test_autoprobe(struct dm_test_state *dms)
 DM_TEST(dm_test_autoprobe, DM_TESTF_SCAN_PDATA);
 
 /* Check that we see the correct platdata in each device */
-static int dm_test_platdata(struct dm_test_state *dms)
+static int dm_test_platdata(struct unit_test_state *uts)
 {
 	const struct dm_test_pdata *pdata;
 	struct udevice *dev;
@@ -270,8 +272,9 @@ static int dm_test_platdata(struct dm_test_state *dms)
 DM_TEST(dm_test_platdata, DM_TESTF_SCAN_PDATA);
 
 /* Test that we can bind, probe, remove, unbind a driver */
-static int dm_test_lifecycle(struct dm_test_state *dms)
+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;
@@ -325,8 +328,9 @@ static int dm_test_lifecycle(struct dm_test_state *dms)
 DM_TEST(dm_test_lifecycle, DM_TESTF_SCAN_PDATA | DM_TESTF_PROBE_TEST);
 
 /* Test that we can bind/unbind and the lists update correctly */
-static int dm_test_ordering(struct dm_test_state *dms)
+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;
 
@@ -380,7 +384,7 @@ static int dm_test_ordering(struct dm_test_state *dms)
 DM_TEST(dm_test_ordering, DM_TESTF_SCAN_PDATA);
 
 /* Check that we can perform operations on a device (do a ping) */
-int dm_check_operations(struct dm_test_state *dms, struct udevice *dev,
+int dm_check_operations(struct unit_test_state *uts, struct udevice *dev,
 			uint32_t base, struct dm_test_priv *priv)
 {
 	int expected;
@@ -408,7 +412,7 @@ int dm_check_operations(struct dm_test_state *dms, struct udevice *dev,
 }
 
 /* Check that we can perform operations on devices */
-static int dm_test_operations(struct dm_test_state *dms)
+static int dm_test_operations(struct unit_test_state *uts)
 {
 	struct udevice *dev;
 	int i;
@@ -430,7 +434,7 @@ static int dm_test_operations(struct dm_test_state *dms)
 		base = test_pdata[i].ping_add;
 		debug("dev=%d, base=%d\n", i, base);
 
-		ut_assert(!dm_check_operations(dms, dev, base, dev->priv));
+		ut_assert(!dm_check_operations(uts, dev, base, dev->priv));
 	}
 
 	return 0;
@@ -438,7 +442,7 @@ static int dm_test_operations(struct dm_test_state *dms)
 DM_TEST(dm_test_operations, DM_TESTF_SCAN_PDATA);
 
 /* Remove all drivers and check that things work */
-static int dm_test_remove(struct dm_test_state *dms)
+static int dm_test_remove(struct unit_test_state *uts)
 {
 	struct udevice *dev;
 	int i;
@@ -460,7 +464,7 @@ static int dm_test_remove(struct dm_test_state *dms)
 DM_TEST(dm_test_remove, DM_TESTF_SCAN_PDATA | DM_TESTF_PROBE_TEST);
 
 /* Remove and recreate everything, check for memory leaks */
-static int dm_test_leak(struct dm_test_state *dms)
+static int dm_test_leak(struct unit_test_state *uts)
 {
 	int i;
 
@@ -469,7 +473,7 @@ static int dm_test_leak(struct dm_test_state *dms)
 		int ret;
 		int id;
 
-		dm_leak_check_start(dms);
+		dm_leak_check_start(uts);
 
 		ut_assertok(dm_scan_platdata(false));
 		ut_assertok(dm_scan_fdt(gd->fdt_blob, false));
@@ -483,7 +487,7 @@ static int dm_test_leak(struct dm_test_state *dms)
 			ut_assertok(ret);
 		}
 
-		ut_assertok(dm_leak_check_end(dms));
+		ut_assertok(dm_leak_check_end(uts));
 	}
 
 	return 0;
@@ -491,7 +495,7 @@ static int dm_test_leak(struct dm_test_state *dms)
 DM_TEST(dm_test_leak, 0);
 
 /* Test uclass init/destroy methods */
-static int dm_test_uclass(struct dm_test_state *dms)
+static int dm_test_uclass(struct unit_test_state *uts)
 {
 	struct uclass *uc;
 
@@ -520,7 +524,7 @@ DM_TEST(dm_test_uclass, 0);
  *		this array.
  * @return 0 if OK, -ve on error
  */
-static int create_children(struct dm_test_state *dms, struct udevice *parent,
+static int create_children(struct unit_test_state *uts, struct udevice *parent,
 			   int count, int key, struct udevice *child[])
 {
 	struct udevice *dev;
@@ -543,8 +547,9 @@ static int create_children(struct dm_test_state *dms, struct udevice *parent,
 
 #define NODE_COUNT	10
 
-static int dm_test_children(struct dm_test_state *dms)
+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];
@@ -559,15 +564,15 @@ static int dm_test_children(struct dm_test_state *dms)
 	ut_assert(NODE_COUNT > 5);
 
 	/* First create 10 top-level children */
-	ut_assertok(create_children(dms, dms->root, NODE_COUNT, 0, top));
+	ut_assertok(create_children(uts, dms->root, NODE_COUNT, 0, top));
 
 	/* Now a few have their own children */
-	ut_assertok(create_children(dms, top[2], NODE_COUNT, 2, NULL));
-	ut_assertok(create_children(dms, top[5], NODE_COUNT, 5, child));
+	ut_assertok(create_children(uts, top[2], NODE_COUNT, 2, NULL));
+	ut_assertok(create_children(uts, top[5], NODE_COUNT, 5, child));
 
 	/* And grandchildren */
 	for (i = 0; i < NODE_COUNT; i++)
-		ut_assertok(create_children(dms, child[i], NODE_COUNT, 50 * i,
+		ut_assertok(create_children(uts, child[i], NODE_COUNT, 50 * i,
 					    i == 2 ? grandchild : NULL));
 
 	/* Check total number of devices */
@@ -629,8 +634,9 @@ static int dm_test_children(struct dm_test_state *dms)
 DM_TEST(dm_test_children, 0);
 
 /* Test that pre-relocation devices work as expected */
-static int dm_test_pre_reloc(struct dm_test_state *dms)
+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 */
@@ -645,7 +651,7 @@ static int dm_test_pre_reloc(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_pre_reloc, 0);
 
-static int dm_test_uclass_before_ready(struct dm_test_state *dms)
+static int dm_test_uclass_before_ready(struct unit_test_state *uts)
 {
 	struct uclass *uc;
 
@@ -661,7 +667,7 @@ static int dm_test_uclass_before_ready(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_uclass_before_ready, 0);
 
-static int dm_test_uclass_devices_find(struct dm_test_state *dms)
+static int dm_test_uclass_devices_find(struct unit_test_state *uts)
 {
 	struct udevice *dev;
 	int ret;
@@ -677,7 +683,7 @@ static int dm_test_uclass_devices_find(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_uclass_devices_find, DM_TESTF_SCAN_PDATA);
 
-static int dm_test_uclass_devices_find_by_name(struct dm_test_state *dms)
+static int dm_test_uclass_devices_find_by_name(struct unit_test_state *uts)
 {
 	struct udevice *finddev;
 	struct udevice *testdev;
@@ -714,7 +720,7 @@ static int dm_test_uclass_devices_find_by_name(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_uclass_devices_find_by_name, DM_TESTF_SCAN_FDT);
 
-static int dm_test_uclass_devices_get(struct dm_test_state *dms)
+static int dm_test_uclass_devices_get(struct unit_test_state *uts)
 {
 	struct udevice *dev;
 	int ret;
@@ -731,7 +737,7 @@ static int dm_test_uclass_devices_get(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_uclass_devices_get, DM_TESTF_SCAN_PDATA);
 
-static int dm_test_uclass_devices_get_by_name(struct dm_test_state *dms)
+static int dm_test_uclass_devices_get_by_name(struct unit_test_state *uts)
 {
 	struct udevice *finddev;
 	struct udevice *testdev;
@@ -775,7 +781,7 @@ static int dm_test_uclass_devices_get_by_name(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_uclass_devices_get_by_name, DM_TESTF_SCAN_FDT);
 
-static int dm_test_device_get_uclass_id(struct dm_test_state *dms)
+static int dm_test_device_get_uclass_id(struct unit_test_state *uts)
 {
 	struct udevice *dev;
 
diff --git a/test/dm/eth.c b/test/dm/eth.c
index f31f6b3..061584e 100644
--- a/test/dm/eth.c
+++ b/test/dm/eth.c
@@ -9,16 +9,16 @@
 
 #include <common.h>
 #include <dm.h>
-#include <dm/test.h>
-#include <dm/ut.h>
 #include <fdtdec.h>
 #include <malloc.h>
 #include <net.h>
+#include <dm/test.h>
 #include <asm/eth.h>
+#include <test/ut.h>
 
 DECLARE_GLOBAL_DATA_PTR;
 
-static int dm_test_eth(struct dm_test_state *dms)
+static int dm_test_eth(struct unit_test_state *uts)
 {
 	net_ping_ip = string_to_ip("1.1.2.2");
 
@@ -38,7 +38,7 @@ static int dm_test_eth(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_eth, DM_TESTF_SCAN_FDT);
 
-static int dm_test_eth_alias(struct dm_test_state *dms)
+static int dm_test_eth_alias(struct unit_test_state *uts)
 {
 	net_ping_ip = string_to_ip("1.1.2.2");
 	setenv("ethact", "eth0");
@@ -62,7 +62,7 @@ static int dm_test_eth_alias(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_eth_alias, DM_TESTF_SCAN_FDT);
 
-static int dm_test_eth_prime(struct dm_test_state *dms)
+static int dm_test_eth_prime(struct unit_test_state *uts)
 {
 	net_ping_ip = string_to_ip("1.1.2.2");
 
@@ -82,7 +82,7 @@ static int dm_test_eth_prime(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_eth_prime, DM_TESTF_SCAN_FDT);
 
-static int dm_test_eth_rotate(struct dm_test_state *dms)
+static int dm_test_eth_rotate(struct unit_test_state *uts)
 {
 	char ethaddr[18];
 
@@ -127,7 +127,7 @@ static int dm_test_eth_rotate(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_eth_rotate, DM_TESTF_SCAN_FDT);
 
-static int dm_test_net_retry(struct dm_test_state *dms)
+static int dm_test_net_retry(struct unit_test_state *uts)
 {
 	net_ping_ip = string_to_ip("1.1.2.2");
 
diff --git a/test/dm/gpio.c b/test/dm/gpio.c
index b29daf1..727db18 100644
--- a/test/dm/gpio.c
+++ b/test/dm/gpio.c
@@ -8,15 +8,15 @@
 #include <fdtdec.h>
 #include <dm.h>
 #include <dm/root.h>
-#include <dm/ut.h>
 #include <dm/test.h>
 #include <dm/util.h>
 #include <asm/gpio.h>
+#include <test/ut.h>
 
 DECLARE_GLOBAL_DATA_PTR;
 
 /* Test that sandbox GPIOs work correctly */
-static int dm_test_gpio(struct dm_test_state *dms)
+static int dm_test_gpio(struct unit_test_state *uts)
 {
 	unsigned int offset, gpio;
 	struct dm_gpio_ops *ops;
@@ -103,7 +103,7 @@ static int dm_test_gpio(struct dm_test_state *dms)
 DM_TEST(dm_test_gpio, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that sandbox anonymous GPIOs work correctly */
-static int dm_test_gpio_anon(struct dm_test_state *dms)
+static int dm_test_gpio_anon(struct unit_test_state *uts)
 {
 	unsigned int offset, gpio;
 	struct udevice *dev;
@@ -125,7 +125,7 @@ static int dm_test_gpio_anon(struct dm_test_state *dms)
 DM_TEST(dm_test_gpio_anon, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that gpio_requestf() works as expected */
-static int dm_test_gpio_requestf(struct dm_test_state *dms)
+static int dm_test_gpio_requestf(struct unit_test_state *uts)
 {
 	unsigned int offset, gpio;
 	struct udevice *dev;
@@ -143,7 +143,7 @@ static int dm_test_gpio_requestf(struct dm_test_state *dms)
 DM_TEST(dm_test_gpio_requestf, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that gpio_request() copies its string */
-static int dm_test_gpio_copy(struct dm_test_state *dms)
+static int dm_test_gpio_copy(struct unit_test_state *uts)
 {
 	unsigned int offset, gpio;
 	struct udevice *dev;
@@ -165,19 +165,19 @@ static int dm_test_gpio_copy(struct dm_test_state *dms)
 DM_TEST(dm_test_gpio_copy, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that we don't leak memory with GPIOs */
-static int dm_test_gpio_leak(struct dm_test_state *dms)
+static int dm_test_gpio_leak(struct unit_test_state *uts)
 {
-	ut_assertok(dm_test_gpio(dms));
-	ut_assertok(dm_test_gpio_anon(dms));
-	ut_assertok(dm_test_gpio_requestf(dms));
-	ut_assertok(dm_leak_check_end(dms));
+	ut_assertok(dm_test_gpio(uts));
+	ut_assertok(dm_test_gpio_anon(uts));
+	ut_assertok(dm_test_gpio_requestf(uts));
+	ut_assertok(dm_leak_check_end(uts));
 
 	return 0;
 }
 DM_TEST(dm_test_gpio_leak, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that we can find GPIOs using phandles */
-static int dm_test_gpio_phandles(struct dm_test_state *dms)
+static int dm_test_gpio_phandles(struct unit_test_state *uts)
 {
 	struct gpio_desc desc, desc_list[8], desc_list2[8];
 	struct udevice *dev, *gpio_a, *gpio_b;
diff --git a/test/dm/i2c.c b/test/dm/i2c.c
index c5939a1..23d612e 100644
--- a/test/dm/i2c.c
+++ b/test/dm/i2c.c
@@ -10,19 +10,19 @@
 #include <dm.h>
 #include <fdtdec.h>
 #include <i2c.h>
+#include <asm/state.h>
+#include <asm/test.h>
 #include <dm/device-internal.h>
 #include <dm/test.h>
 #include <dm/uclass-internal.h>
-#include <dm/ut.h>
 #include <dm/util.h>
-#include <asm/state.h>
-#include <asm/test.h>
+#include <test/ut.h>
 
 static const int busnum;
 static const int chip = 0x2c;
 
 /* Test that we can find buses and chips */
-static int dm_test_i2c_find(struct dm_test_state *dms)
+static int dm_test_i2c_find(struct unit_test_state *uts)
 {
 	struct udevice *bus, *dev;
 	const int no_chip = 0x10;
@@ -43,7 +43,7 @@ static int dm_test_i2c_find(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_i2c_find, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
-static int dm_test_i2c_read_write(struct dm_test_state *dms)
+static int dm_test_i2c_read_write(struct unit_test_state *uts)
 {
 	struct udevice *bus, *dev;
 	uint8_t buf[5];
@@ -60,7 +60,7 @@ static int dm_test_i2c_read_write(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_i2c_read_write, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
-static int dm_test_i2c_speed(struct dm_test_state *dms)
+static int dm_test_i2c_speed(struct unit_test_state *uts)
 {
 	struct udevice *bus, *dev;
 	uint8_t buf[5];
@@ -82,7 +82,7 @@ static int dm_test_i2c_speed(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_i2c_speed, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
-static int dm_test_i2c_offset_len(struct dm_test_state *dms)
+static int dm_test_i2c_offset_len(struct unit_test_state *uts)
 {
 	struct udevice *bus, *dev;
 	uint8_t buf[5];
@@ -99,7 +99,7 @@ static int dm_test_i2c_offset_len(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_i2c_offset_len, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
-static int dm_test_i2c_probe_empty(struct dm_test_state *dms)
+static int dm_test_i2c_probe_empty(struct unit_test_state *uts)
 {
 	struct udevice *bus, *dev;
 
@@ -114,7 +114,7 @@ static int dm_test_i2c_probe_empty(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_i2c_probe_empty, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
-static int dm_test_i2c_bytewise(struct dm_test_state *dms)
+static int dm_test_i2c_bytewise(struct unit_test_state *uts)
 {
 	struct udevice *bus, *dev;
 	struct udevice *eeprom;
@@ -169,7 +169,7 @@ static int dm_test_i2c_bytewise(struct dm_test_state *dms)
 }
 DM_TEST(dm_test_i2c_bytewise, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
-static int dm_test_i2c_offset(struct dm_test_state *dms)
+static int dm_test_i2c_offset(struct unit_test_state *uts)
 {
 	struct udevice *eeprom;
 	struct udevice *dev;
diff --git a/test/dm/pci.c b/test/dm/pci.c
index 6c63fa4..2f3ae79 100644
--- a/test/dm/pci.c
+++ b/test/dm/pci.c
@@ -8,10 +8,10 @@
 #include <dm.h>
 #include <asm/io.h>
 #include <dm/test.h>
-#include <dm/ut.h>
+#include <test/ut.h>
 
 /* Test that sandbox PCI works correctly */
-static int dm_test_pci_base(struct dm_test_state *dms)
+static int dm_test_pci_base(struct unit_test_state *uts)
 {
 	struct udevice *bus;
 
@@ -22,7 +22,7 @@ static int dm_test_pci_base(struct dm_test_state *dms)
 DM_TEST(dm_test_pci_base, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that we can use the swapcase device correctly */
-static int dm_test_pci_swapcase(struct dm_test_state *dms)
+static int dm_test_pci_swapcase(struct unit_test_state *uts)
 {
 	pci_dev_t pci_dev = PCI_BDF(0, 0x1f, 0);
 	struct pci_controller *hose;
diff --git a/test/dm/pmic.c b/test/dm/pmic.c
index e9c904c..422ea3e 100644
--- a/test/dm/pmic.c
+++ b/test/dm/pmic.c
@@ -14,17 +14,17 @@
 #include <malloc.h>
 #include <dm/device-internal.h>
 #include <dm/root.h>
-#include <dm/ut.h>
 #include <dm/util.h>
 #include <dm/test.h>
 #include <dm/uclass-internal.h>
 #include <power/pmic.h>
 #include <power/sandbox_pmic.h>
+#include <test/ut.h>
 
 DECLARE_GLOBAL_DATA_PTR;
 
 /* Test PMIC get method */
-static int dm_test_power_pmic_get(struct dm_test_state *dms)
+static int dm_test_power_pmic_get(struct unit_test_state *uts)
 {
 	const char *name = "sandbox_pmic";
 	struct udevice *dev;
@@ -40,7 +40,7 @@ static int dm_test_power_pmic_get(struct dm_test_state *dms)
 DM_TEST(dm_test_power_pmic_get, DM_TESTF_SCAN_FDT);
 
 /* Test PMIC I/O */
-static int dm_test_power_pmic_io(struct dm_test_state *dms)
+static int dm_test_power_pmic_io(struct unit_test_state *uts)
 {
 	const char *name = "sandbox_pmic";
 	uint8_t out_buffer, in_buffer;
diff --git a/test/dm/regulator.c b/test/dm/regulator.c
index c4f14bd..d279c04 100644
--- a/test/dm/regulator.c
+++ b/test/dm/regulator.c
@@ -14,13 +14,13 @@
 #include <malloc.h>
 #include <dm/device-internal.h>
 #include <dm/root.h>
-#include <dm/ut.h>
 #include <dm/util.h>
 #include <dm/test.h>
 #include <dm/uclass-internal.h>
 #include <power/pmic.h>
 #include <power/regulator.h>
 #include <power/sandbox_pmic.h>
+#include <test/ut.h>
 
 DECLARE_GLOBAL_DATA_PTR;
 
@@ -47,7 +47,7 @@ static const char *regulator_names[OUTPUT_COUNT][OUTPUT_NAME_COUNT] = {
 };
 
 /* Test regulator get method */
-static int dm_test_power_regulator_get(struct dm_test_state *dms)
+static int dm_test_power_regulator_get(struct unit_test_state *uts)
 {
 	struct dm_regulator_uclass_platdata *uc_pdata;
 	struct udevice *dev_by_devname;
@@ -92,7 +92,7 @@ static int dm_test_power_regulator_get(struct dm_test_state *dms)
 DM_TEST(dm_test_power_regulator_get, DM_TESTF_SCAN_FDT);
 
 /* Test regulator set and get Voltage method */
-static int dm_test_power_regulator_set_get_voltage(struct dm_test_state *dms)
+static int dm_test_power_regulator_set_get_voltage(struct unit_test_state *uts)
 {
 	struct dm_regulator_uclass_platdata *uc_pdata;
 	struct udevice *dev;
@@ -119,7 +119,7 @@ static int dm_test_power_regulator_set_get_voltage(struct dm_test_state *dms)
 DM_TEST(dm_test_power_regulator_set_get_voltage, DM_TESTF_SCAN_FDT);
 
 /* Test regulator set and get Current method */
-static int dm_test_power_regulator_set_get_current(struct dm_test_state *dms)
+static int dm_test_power_regulator_set_get_current(struct unit_test_state *uts)
 {
 	struct dm_regulator_uclass_platdata *uc_pdata;
 	struct udevice *dev;
@@ -158,7 +158,7 @@ static int dm_test_power_regulator_set_get_current(struct dm_test_state *dms)
 DM_TEST(dm_test_power_regulator_set_get_current, DM_TESTF_SCAN_FDT);
 
 /* Test regulator set and get Enable method */
-static int dm_test_power_regulator_set_get_enable(struct dm_test_state *dms)
+static int dm_test_power_regulator_set_get_enable(struct unit_test_state *uts)
 {
 	const char *platname;
 	struct udevice *dev;
@@ -177,7 +177,7 @@ static int dm_test_power_regulator_set_get_enable(struct dm_test_state *dms)
 DM_TEST(dm_test_power_regulator_set_get_enable, DM_TESTF_SCAN_FDT);
 
 /* Test regulator set and get mode method */
-static int dm_test_power_regulator_set_get_mode(struct dm_test_state *dms)
+static int dm_test_power_regulator_set_get_mode(struct unit_test_state *uts)
 {
 	const char *platname;
 	struct udevice *dev;
@@ -196,7 +196,7 @@ static int dm_test_power_regulator_set_get_mode(struct dm_test_state *dms)
 DM_TEST(dm_test_power_regulator_set_get_mode, DM_TESTF_SCAN_FDT);
 
 /* Test regulator autoset method */
-static int dm_test_power_regulator_autoset(struct dm_test_state *dms)
+static int dm_test_power_regulator_autoset(struct unit_test_state *uts)
 {
 	const char *platname;
 	struct udevice *dev, *dev_autoset;
@@ -276,7 +276,7 @@ static const struct setting expected_setting_list[] = {
 static int list_count = ARRAY_SIZE(expected_setting_list);
 
 /* Test regulator list autoset method */
-static int dm_test_power_regulator_autoset_list(struct dm_test_state *dms)
+static int dm_test_power_regulator_autoset_list(struct unit_test_state *uts)
 {
 	struct udevice *dev_list[2], *dev;
 	int i;
diff --git a/test/dm/rtc.c b/test/dm/rtc.c
index 9397cf7..4345708 100644
--- a/test/dm/rtc.c
+++ b/test/dm/rtc.c
@@ -9,12 +9,12 @@
 #include <dm.h>
 #include <rtc.h>
 #include <asm/io.h>
-#include <dm/test.h>
-#include <dm/ut.h>
 #include <asm/test.h>
+#include <dm/test.h>
+#include <test/ut.h>
 
 /* Simple RTC sanity check */
-static int dm_test_rtc_base(struct dm_test_state *dms)
+static int dm_test_rtc_base(struct unit_test_state *uts)
 {
 	struct udevice *dev;
 
@@ -52,7 +52,7 @@ static int cmp_times(struct rtc_time *expect, struct rtc_time *time, bool show)
 }
 
 /* Set and get the time */
-static int dm_test_rtc_set_get(struct dm_test_state *dms)
+static int dm_test_rtc_set_get(struct unit_test_state *uts)
 {
 	struct rtc_time now, time, cmp;
 	struct udevice *dev, *emul;
@@ -117,7 +117,7 @@ static int dm_test_rtc_set_get(struct dm_test_state *dms)
 DM_TEST(dm_test_rtc_set_get, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Reset the time */
-static int dm_test_rtc_reset(struct dm_test_state *dms)
+static int dm_test_rtc_reset(struct unit_test_state *uts)
 {
 	struct rtc_time now;
 	struct udevice *dev, *emul;
@@ -143,7 +143,7 @@ static int dm_test_rtc_reset(struct dm_test_state *dms)
 DM_TEST(dm_test_rtc_reset, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Check that two RTC devices can be used independently */
-static int dm_test_rtc_dual(struct dm_test_state *dms)
+static int dm_test_rtc_dual(struct unit_test_state *uts)
 {
 	struct rtc_time now1, now2, cmp;
 	struct udevice *dev1, *dev2;
diff --git a/test/dm/sf.c b/test/dm/sf.c
index 08098a1..b084462 100644
--- a/test/dm/sf.c
+++ b/test/dm/sf.c
@@ -10,12 +10,12 @@
 #include <spi.h>
 #include <spi_flash.h>
 #include <asm/state.h>
-#include <dm/ut.h>
 #include <dm/test.h>
 #include <dm/util.h>
+#include <test/ut.h>
 
 /* Test that sandbox SPI flash works correctly */
-static int dm_test_spi_flash(struct dm_test_state *dms)
+static int dm_test_spi_flash(struct unit_test_state *uts)
 {
 	/*
 	 * Create an empty test file and run the SPI flash tests. This is a
diff --git a/test/dm/spi.c b/test/dm/spi.c
index c7ee652..2e27da7 100644
--- a/test/dm/spi.c
+++ b/test/dm/spi.c
@@ -9,15 +9,15 @@
 #include <fdtdec.h>
 #include <spi.h>
 #include <spi_flash.h>
+#include <asm/state.h>
 #include <dm/device-internal.h>
 #include <dm/test.h>
 #include <dm/uclass-internal.h>
-#include <dm/ut.h>
 #include <dm/util.h>
-#include <asm/state.h>
+#include <test/ut.h>
 
 /* Test that we can find buses and chip-selects */
-static int dm_test_spi_find(struct dm_test_state *dms)
+static int dm_test_spi_find(struct unit_test_state *uts)
 {
 	struct sandbox_state *state = state_get_current();
 	struct spi_slave *slave;
@@ -95,7 +95,7 @@ static int dm_test_spi_find(struct dm_test_state *dms)
 DM_TEST(dm_test_spi_find, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that sandbox SPI works correctly */
-static int dm_test_spi_xfer(struct dm_test_state *dms)
+static int dm_test_spi_xfer(struct unit_test_state *uts)
 {
 	struct spi_slave *slave;
 	struct udevice *bus;
diff --git a/test/dm/test-driver.c b/test/dm/test-driver.c
index bc6a6e7..d10af51 100644
--- a/test/dm/test-driver.c
+++ b/test/dm/test-driver.c
@@ -12,11 +12,11 @@
 #include <errno.h>
 #include <malloc.h>
 #include <dm/test.h>
-#include <dm/ut.h>
+#include <test/ut.h>
 #include <asm/io.h>
 
 int dm_testdrv_op_count[DM_TEST_OP_COUNT];
-static struct dm_test_state *dms = &global_test_state;
+static struct unit_test_state *uts = &global_dm_test_state;
 
 static int testdrv_ping(struct udevice *dev, int pingval, int *pingret)
 {
@@ -114,6 +114,8 @@ static int test_manual_bind(struct udevice *dev)
 
 static int test_manual_probe(struct udevice *dev)
 {
+	struct dm_test_state *dms = uts->priv;
+
 	dm_testdrv_op_count[DM_TEST_OP_PROBE]++;
 	if (!dms->force_fail_alloc)
 		dev->priv = calloc(1, sizeof(struct dm_test_priv));
diff --git a/test/dm/test-fdt.c b/test/dm/test-fdt.c
index b8ee959..49a36cb 100644
--- a/test/dm/test-fdt.c
+++ b/test/dm/test-fdt.c
@@ -12,9 +12,9 @@
 #include <asm/io.h>
 #include <dm/test.h>
 #include <dm/root.h>
-#include <dm/ut.h>
 #include <dm/uclass-internal.h>
 #include <dm/util.h>
+#include <test/ut.h>
 
 DECLARE_GLOBAL_DATA_PTR;
 
@@ -99,7 +99,7 @@ UCLASS_DRIVER(testfdt) = {
 	.flags		= DM_UC_FLAG_SEQ_ALIAS,
 };
 
-int dm_check_devices(struct dm_test_state *dms, int num_devices)
+int dm_check_devices(struct unit_test_state *uts, int num_devices)
 {
 	struct udevice *dev;
 	int ret;
@@ -126,7 +126,7 @@ int dm_check_devices(struct dm_test_state *dms, int num_devices)
 		debug("dev=%d, base=%d: %s\n", i, base,
 		      fdt_get_name(gd->fdt_blob, dev->of_offset, NULL));
 
-		ut_assert(!dm_check_operations(dms, dev, base,
+		ut_assert(!dm_check_operations(uts, dev, base,
 					       dev_get_priv(dev)));
 	}
 
@@ -134,7 +134,7 @@ int dm_check_devices(struct dm_test_state *dms, int num_devices)
 }
 
 /* Test that FDT-based binding works correctly */
-static int dm_test_fdt(struct dm_test_state *dms)
+static int dm_test_fdt(struct unit_test_state *uts)
 {
 	const int num_devices = 6;
 	struct udevice *dev;
@@ -159,13 +159,13 @@ static int dm_test_fdt(struct dm_test_state *dms)
 		ut_assert(dev->platdata);
 	}
 
-	ut_assertok(dm_check_devices(dms, num_devices));
+	ut_assertok(dm_check_devices(uts, num_devices));
 
 	return 0;
 }
 DM_TEST(dm_test_fdt, 0);
 
-static int dm_test_fdt_pre_reloc(struct dm_test_state *dms)
+static int dm_test_fdt_pre_reloc(struct unit_test_state *uts)
 {
 	struct uclass *uc;
 	int ret;
@@ -184,7 +184,7 @@ static int dm_test_fdt_pre_reloc(struct dm_test_state *dms)
 DM_TEST(dm_test_fdt_pre_reloc, 0);
 
 /* Test that sequence numbers are allocated properly */
-static int dm_test_fdt_uclass_seq(struct dm_test_state *dms)
+static int dm_test_fdt_uclass_seq(struct unit_test_state *uts)
 {
 	struct udevice *dev;
 
@@ -239,7 +239,7 @@ static int dm_test_fdt_uclass_seq(struct dm_test_state *dms)
 DM_TEST(dm_test_fdt_uclass_seq, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
 
 /* Test that we can find a device by device tree offset */
-static int dm_test_fdt_offset(struct dm_test_state *dms)
+static int dm_test_fdt_offset(struct unit_test_state *uts)
 {
 	const void *blob = gd->fdt_blob;
 	struct udevice *dev;
diff --git a/test/dm/test-main.c b/test/dm/test-main.c
index 7348f69..5e36e76 100644
--- a/test/dm/test-main.c
+++ b/test/dm/test-main.c
@@ -11,15 +11,20 @@
 #include <dm/test.h>
 #include <dm/root.h>
 #include <dm/uclass-internal.h>
-#include <dm/ut.h>
+#include <test/ut.h>
 
 DECLARE_GLOBAL_DATA_PTR;
 
-struct dm_test_state global_test_state;
+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 dm_test_state *dms)
+static int dm_test_init(struct unit_test_state *uts)
 {
+	struct dm_test_state *dms = uts->priv;
+
+	memset(uts, '\0', sizeof(*uts));
+	uts->priv = dms;
 	memset(dms, '\0', sizeof(*dms));
 	gd->dm_root = NULL;
 	memset(dm_testdrv_op_count, '\0', sizeof(dm_testdrv_op_count));
@@ -31,7 +36,7 @@ static int dm_test_init(struct dm_test_state *dms)
 }
 
 /* Ensure all the test devices are probed */
-static int do_autoprobe(struct dm_test_state *dms)
+static int do_autoprobe(struct unit_test_state *uts)
 {
 	struct udevice *dev;
 	int ret;
@@ -45,7 +50,7 @@ static int do_autoprobe(struct dm_test_state *dms)
 	return ret;
 }
 
-static int dm_test_destroy(struct dm_test_state *dms)
+static int dm_test_destroy(struct unit_test_state *uts)
 {
 	int id;
 
@@ -67,10 +72,11 @@ static int dm_test_destroy(struct dm_test_state *dms)
 
 int dm_test_main(const char *test_name)
 {
-	struct dm_test *tests = ll_entry_start(struct dm_test, dm_test);
-	const int n_ents = ll_entry_count(struct dm_test, dm_test);
-	struct dm_test_state *dms = &global_test_state;
-	struct dm_test *test;
+	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;
+	uts->priv = &_global_priv_dm_test_state;
+	struct unit_test *test;
 
 	/*
 	 * If we have no device tree, or it only has a root node, then these
@@ -89,23 +95,23 @@ int dm_test_main(const char *test_name)
 		if (test_name && strcmp(test_name, test->name))
 			continue;
 		printf("Test: %s\n", test->name);
-		ut_assertok(dm_test_init(dms));
+		ut_assertok(dm_test_init(uts));
 
-		dms->start = mallinfo();
+		uts->start = mallinfo();
 		if (test->flags & DM_TESTF_SCAN_PDATA)
 			ut_assertok(dm_scan_platdata(false));
 		if (test->flags & DM_TESTF_PROBE_TEST)
-			ut_assertok(do_autoprobe(dms));
+			ut_assertok(do_autoprobe(uts));
 		if (test->flags & DM_TESTF_SCAN_FDT)
 			ut_assertok(dm_scan_fdt(gd->fdt_blob, false));
 
-		if (test->func(dms))
+		if (test->func(uts))
 			break;
 
-		ut_assertok(dm_test_destroy(dms));
+		ut_assertok(dm_test_destroy(uts));
 	}
 
-	printf("Failures: %d\n", dms->fail_count);
+	printf("Failures: %d\n", uts->fail_count);
 
 	return 0;
 }
diff --git a/test/dm/test-uclass.c b/test/dm/test-uclass.c
index 4ae75ef..4a543bb 100644
--- a/test/dm/test-uclass.c
+++ b/test/dm/test-uclass.c
@@ -11,12 +11,12 @@
 #include <malloc.h>
 #include <dm.h>
 #include <errno.h>
-#include <dm/test.h>
-#include <dm/ut.h>
 #include <asm/io.h>
+#include <dm/test.h>
 #include <linux/list.h>
+#include <test/ut.h>
 
-static struct dm_test_state *dms = &global_test_state;
+static struct unit_test_state *uts = &global_dm_test_state;
 
 int test_ping(struct udevice *dev, int pingval, int *pingret)
 {
@@ -70,6 +70,7 @@ static int test_post_probe(struct udevice *dev)
 
 	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);
diff --git a/test/dm/usb.c b/test/dm/usb.c
index 6ea86d7..9939d83 100644
--- a/test/dm/usb.c
+++ b/test/dm/usb.c
@@ -9,10 +9,10 @@
 #include <usb.h>
 #include <asm/io.h>
 #include <dm/test.h>
-#include <dm/ut.h>
+#include <test/ut.h>
 
 /* Test that sandbox USB works correctly */
-static int dm_test_usb_base(struct dm_test_state *dms)
+static int dm_test_usb_base(struct unit_test_state *uts)
 {
 	struct udevice *bus;
 
@@ -29,7 +29,7 @@ DM_TEST(dm_test_usb_base, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
  * covers scanning the bug, setting up a hub and a flash stick and reading
  * data from the flash stick.
  */
-static int dm_test_usb_flash(struct dm_test_state *dms)
+static int dm_test_usb_flash(struct unit_test_state *uts)
 {
 	struct udevice *dev;
 	block_dev_desc_t *dev_desc;
diff --git a/test/dm/ut.c b/test/ut.c
index 8b69bc2..0282de5 100644
--- a/test/dm/ut.c
+++ b/test/ut.c
@@ -1,5 +1,5 @@
 /*
- * Simple unit test library for driver model
+ * Simple unit test library
  *
  * Copyright (c) 2013 Google, Inc
  *
@@ -7,19 +7,17 @@
  */
 
 #include <common.h>
-#include <dm/test.h>
-#include <dm/ut.h>
+#include <test/test.h>
+#include <test/ut.h>
 
-struct dm_test_state;
-
-void ut_fail(struct dm_test_state *dms, const char *fname, int line,
+void ut_fail(struct unit_test_state *uts, const char *fname, int line,
 	     const char *func, const char *cond)
 {
 	printf("%s:%d, %s(): %s\n", fname, line, func, cond);
-	dms->fail_count++;
+	uts->fail_count++;
 }
 
-void ut_failf(struct dm_test_state *dms, const char *fname, int line,
+void ut_failf(struct unit_test_state *uts, const char *fname, int line,
 	      const char *func, const char *cond, const char *fmt, ...)
 {
 	va_list args;
@@ -29,5 +27,5 @@ void ut_failf(struct dm_test_state *dms, const char *fname, int line,
 	vprintf(fmt, args);
 	va_end(args);
 	putc('\n');
-	dms->fail_count++;
+	uts->fail_count++;
 }