diff --git a/include/dm/test.h b/include/dm/test.h
index 235d728..f08c05d 100644
--- a/include/dm/test.h
+++ b/include/dm/test.h
@@ -8,6 +8,7 @@
 #define __DM_TEST_H
 
 #include <dm.h>
+#include <malloc.h>
 
 /**
  * struct dm_test_cdata - configuration data for test instance
@@ -120,6 +121,7 @@ struct dm_test_state {
 	int force_fail_alloc;
 	int skip_post_probe;
 	struct udevice *removed;
+	struct mallinfo start;
 };
 
 /* Test flags for each test */
@@ -178,6 +180,27 @@ int dm_check_operations(struct dm_test_state *dms, struct udevice *dev,
 int dm_check_devices(struct dm_test_state *dms, int num_devices);
 
 /**
+ * dm_leak_check_start() - Prepare to check for a memory leak
+ *
+ * Call this before allocating memory to record the amount of memory being
+ * used.
+ *
+ * @dms: Overall test state
+ */
+void dm_leak_check_start(struct dm_test_state *dms);
+
+/**
+ * dm_leak_check_end() - Check that no memory has leaked
+ *
+ * Call this after dm_leak_check_start() and after you have hopefuilly freed
+ * all the memory that was allocated. This function will print an error if
+ * 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);
+
+
+/**
  * dm_test_main() - Run all the tests
  *
  * This runs all available driver model tests
diff --git a/test/dm/core.c b/test/dm/core.c
index b0cfb42..ff5c2a7 100644
--- a/test/dm/core.c
+++ b/test/dm/core.c
@@ -67,6 +67,34 @@ static struct driver_info driver_info_pre_reloc = {
 	.platdata = &test_pdata_manual,
 };
 
+void dm_leak_check_start(struct dm_test_state *dms)
+{
+	dms->start = mallinfo();
+	if (!dms->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)
+{
+	struct mallinfo end;
+	int id;
+
+	/* Don't delete the root class, since we started with that */
+	for (id = UCLASS_ROOT + 1; id < UCLASS_COUNT; id++) {
+		struct uclass *uc;
+
+		uc = uclass_find(id);
+		if (!uc)
+			continue;
+		ut_assertok(uclass_destroy(uc));
+	}
+
+	end = mallinfo();
+	ut_asserteq(dms->start.uordblks, end.uordblks);
+
+	return 0;
+}
+
 /* Test that binding with platdata occurs correctly */
 static int dm_test_autobind(struct dm_test_state *dms)
 {
@@ -377,14 +405,11 @@ static int dm_test_leak(struct dm_test_state *dms)
 	int i;
 
 	for (i = 0; i < 2; i++) {
-		struct mallinfo start, end;
 		struct udevice *dev;
 		int ret;
 		int id;
 
-		start = mallinfo();
-		if (!start.uordblks)
-			puts("Warning: Please add '#define DEBUG' to the top of common/dlmalloc.c\n");
+		dm_leak_check_start(dms);
 
 		ut_assertok(dm_scan_platdata(false));
 		ut_assertok(dm_scan_fdt(gd->fdt_blob, false));
@@ -398,18 +423,7 @@ static int dm_test_leak(struct dm_test_state *dms)
 			ut_assertok(ret);
 		}
 
-		/* Don't delete the root class, since we started with that */
-		for (id = UCLASS_ROOT + 1; id < UCLASS_COUNT; id++) {
-			struct uclass *uc;
-
-			uc = uclass_find(id);
-			if (!uc)
-				continue;
-			ut_assertok(uclass_destroy(uc));
-		}
-
-		end = mallinfo();
-		ut_asserteq(start.uordblks, end.uordblks);
+		ut_assertok(dm_leak_check_end(dms));
 	}
 
 	return 0;
diff --git a/test/dm/test-main.c b/test/dm/test-main.c
index 94ce72a..90ca810 100644
--- a/test/dm/test-main.c
+++ b/test/dm/test-main.c
@@ -7,6 +7,7 @@
 #include <common.h>
 #include <dm.h>
 #include <errno.h>
+#include <malloc.h>
 #include <dm/test.h>
 #include <dm/root.h>
 #include <dm/uclass-internal.h>
@@ -88,6 +89,7 @@ int dm_test_main(void)
 		printf("Test: %s\n", test->name);
 		ut_assertok(dm_test_init(dms));
 
+		dms->start = mallinfo();
 		if (test->flags & DM_TESTF_SCAN_PDATA)
 			ut_assertok(dm_scan_platdata(false));
 		if (test->flags & DM_TESTF_PROBE_TEST)