From 88a1ae8172a847911ce699ac95f79ea17c88739c Mon Sep 17 00:00:00 2001
From: Simon Glass <sjg@chromium.org>
Date: Tue, 6 Sep 2022 20:27:29 -0600
Subject: [PATCH] dm: core: Create a function to get a live tree in a test

Move this logic out of the test into separate functions, so we can use it
in other tests.

Signed-off-by: Simon Glass <sjg@chromium.org>
---
 test/dm/ofnode.c | 57 +++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 49 insertions(+), 8 deletions(-)

diff --git a/test/dm/ofnode.c b/test/dm/ofnode.c
index 134a307f1b..be5e3155e3 100644
--- a/test/dm/ofnode.c
+++ b/test/dm/ofnode.c
@@ -51,6 +51,46 @@ oftree get_other_oftree(struct unit_test_state *uts)
 	return tree;
 }
 
+/**
+ * get_oftree() - Convert a flat tree into an oftree object
+ *
+ * @uts: Test state
+ * @fdt: Pointer to flat tree
+ * @treep: Returns the tree, on success
+ * Return: 0 if OK, 1 if the tree failed to unflatten, -EOVERFLOW if there are
+ * too many flat trees to allow another one to be registers (see
+ * oftree_ensure())
+ */
+int get_oftree(struct unit_test_state *uts, void *fdt, oftree *treep)
+{
+	oftree tree;
+
+	if (of_live_active()) {
+		struct device_node *root;
+
+		ut_assertok(unflatten_device_tree(fdt, &root));
+		tree = oftree_from_np(root);
+	} else {
+		tree = oftree_from_fdt(fdt);
+		if (!oftree_valid(tree))
+			return -EOVERFLOW;
+	}
+	*treep = tree;
+
+	return 0;
+}
+
+/**
+ * free_oftree() - Free memory used by get_oftree()
+ *
+ * @tree: Tree to free
+ */
+void free_oftree(oftree tree)
+{
+	if (of_live_active())
+		free(tree.np);
+}
+
 static int dm_test_ofnode_compatible(struct unit_test_state *uts)
 {
 	ofnode root_node = ofnode_path("/");
@@ -590,10 +630,10 @@ static int make_ofnode_fdt(struct unit_test_state *uts, void *fdt, int size)
 
 static int dm_test_ofnode_root(struct unit_test_state *uts)
 {
-	struct device_node *root = NULL;
 	char fdt[256];
 	oftree tree;
 	ofnode node;
+	int ret;
 
 	/* Check that aliases work on the control FDT */
 	node = ofnode_get_aliases_node("ethernet3");
@@ -603,12 +643,13 @@ static int dm_test_ofnode_root(struct unit_test_state *uts)
 	ut_assert(!oftree_valid(oftree_null()));
 
 	ut_assertok(make_ofnode_fdt(uts, fdt, sizeof(fdt)));
-	if (of_live_active()) {
-		ut_assertok(unflatten_device_tree(fdt, &root));
-		tree = oftree_from_np(root);
-	} else {
-		tree = oftree_from_fdt(fdt);
-	}
+	ret = get_oftree(uts, fdt, &tree);
+
+	/* skip the rest of this test if multiple FDTs are not supported */
+	if (ret == -EOVERFLOW)
+		return 0;
+
+	ut_assertok(ret);
 	ut_assert(oftree_valid(tree));
 
 	/* Make sure they don't work on this new tree */
@@ -623,7 +664,7 @@ static int dm_test_ofnode_root(struct unit_test_state *uts)
 	node = oftree_path(oftree_default(), "/new-mmc");
 	ut_assert(!ofnode_valid(node));
 
-	free(root);
+	free_oftree(tree);
 
 	return 0;
 }
-- 
2.39.5