From c908ecb7b51e886469c4fc8eb5492bdbda55b871 Mon Sep 17 00:00:00 2001
From: Simon Glass <sjg@chromium.org>
Date: Fri, 7 Feb 2025 11:30:38 -0700
Subject: [PATCH] test: Support an init/uninit functions for test suites

Some suites need things to be set up before they can run. Add a way to
declare an init function using the UNIT_TEST_INIT() macro. The init
function is just like any other test, but is always placed first so that
it runs before all the other test functions in the suite.

Add an uninit function as well, to clean up after the test.

Signed-off-by: Simon Glass <sjg@chromium.org>
---
 include/test/test.h | 20 ++++++++++++++++++++
 test/test-main.c    |  3 ++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/include/test/test.h b/include/test/test.h
index 877fda8d5aa..0f2b68a5dee 100644
--- a/include/test/test.h
+++ b/include/test/test.h
@@ -100,6 +100,8 @@ enum ut_flags {
 	UTF_ETH_BOOTDEV	= BIT(9),	/* enable Ethernet bootdevs */
 	UTF_SF_BOOTDEV	= BIT(10),	/* enable SPI flash bootdevs */
 	UFT_BLOBLIST	= BIT(11),	/* test changes gd->bloblist */
+	UTF_INIT	= BIT(12),	/* test inits a suite */
+	UTF_UNINIT	= BIT(13),	/* test uninits a suite */
 };
 
 /**
@@ -147,6 +149,24 @@ struct unit_test {
 		.func = _name,						\
 	}
 
+/* init function for unit-test suite (the 'A' makes it first) */
+#define UNIT_TEST_INIT(_name, _flags, _suite)				\
+	ll_entry_declare(struct unit_test, A ## _name, ut_ ## _suite) = {	\
+		.file = __FILE__,					\
+		.name = #_name,						\
+		.flags = (_flags) | UTF_INIT,				\
+		.func = _name,						\
+	}
+
+/* uninit function for unit-test suite (the 'aaa' makes it last) */
+#define UNIT_TEST_UNINIT(_name, _flags, _suite)				\
+	ll_entry_declare(struct unit_test, zzz ## _name, ut_ ## _suite) = { \
+		.file = __FILE__,					\
+		.name = #_name,						\
+		.flags = (_flags) | UTF_UNINIT,				\
+		.func = _name,						\
+	}
+
 /* Get the start of a list of unit tests for a particular suite */
 #define UNIT_TEST_SUITE_START(_suite) \
 	ll_entry_start(struct unit_test, ut_ ## _suite)
diff --git a/test/test-main.c b/test/test-main.c
index 8c0d820032c..ee855bfe2ed 100644
--- a/test/test-main.c
+++ b/test/test-main.c
@@ -620,7 +620,8 @@ static int ut_run_tests(struct unit_test_state *uts, const char *prefix,
 		const char *test_name = test->name;
 		int ret, i, old_fail_count;
 
-		if (!test_matches(prefix, test_name, select_name))
+		if (!(test->flags & (UTF_INIT | UTF_UNINIT)) &&
+		    !test_matches(prefix, test_name, select_name))
 			continue;
 
 		if (test->flags & UTF_MANUAL) {