test: Move restoring of driver model state to ut_run_list()

Add this functionality to ut_run_list() so it can be removed from
dm_test_run().

At this point all tests are run through ut_run_list().

Signed-off-by: Simon Glass <sjg@chromium.org>

2 files changed, 29 insertions, 12 deletions

diff --git a/test/dm/test-dm.c b/test/dm/test-dm.c
index 8cb99ed80c..cb4f99537d 100644
--- a/test/dm/test-dm.c
+++ b/test/dm/test-dm.c
@@ -24,21 +24,10 @@ int dm_test_run(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 device_node *of_root;
 	int ret;
 
-	of_root = gd_of_root();
 	ret = ut_run_list("driver model", "dm_test_", tests, n_ents, test_name);
 
-	/* Put everything back to normal so that sandbox works as expected */
-	gd_set_of_root(of_root);
-	gd->dm_root = NULL;
-	if (dm_init(CONFIG_IS_ENABLED(OF_LIVE)))
-		return CMD_RET_FAILURE;
-	dm_scan_plat(false);
-	if (!CONFIG_IS_ENABLED(OF_PLATDATA))
-		dm_scan_fdt(false);
-
 	return ret ? CMD_RET_FAILURE : 0;
 }
 
diff --git a/test/test-main.c b/test/test-main.c
index 8138fb4387..6edd49f0b6 100644
--- a/test/test-main.c
+++ b/test/test-main.c
@@ -143,7 +143,7 @@ static bool test_matches(const char *prefix, const char *test_name,
 	return false;
 }
 
-/*
+/**
  * ut_list_has_dm_tests() - Check if a list of tests has driver model ones
  *
  * @tests: List of tests to run
@@ -163,6 +163,28 @@ static bool ut_list_has_dm_tests(struct unit_test *tests, int count)
 }
 
 /**
+ * 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
@@ -359,10 +381,12 @@ 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...
@@ -385,5 +409,9 @@ int ut_run_list(const char *category, const char *prefix,
 	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;
 }