[PATCH 2/9] kunit: Introduce param_init/exit for parameterized test shared context management

[Marie Zhussupova <marievic@google.com>]

Add `param_init` and `param_exit` function pointers to
`struct kunit_case`. Users will be able to set them
via the new `KUNIT_CASE_PARAM_WITH_INIT` macro.

These functions are invoked by kunit_run_tests() once before
and once after the entire parameterized test series, respectively.
They will receive the parent kunit test instance, allowing users
to register and manage shared resources. Resources added to this
parent kunit test will be accessible to all individual parameterized
tests, facilitating init and exit for shared state.

Signed-off-by: Marie Zhussupova <marievic@google.com>
---
 include/kunit/test.h | 33 ++++++++++++++++++++++++++++++++-
 lib/kunit/test.c     | 23 ++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/include/kunit/test.h b/include/kunit/test.h
index a42d0c8cb985..d8dac7efd745 100644
--- a/include/kunit/test.h
+++ b/include/kunit/test.h
@@ -92,6 +92,8 @@ struct kunit_attributes {
  * @name:     the name of the test case.
  * @generate_params: the generator function for parameterized tests.
  * @attr:     the attributes associated with the test
+ * @param_init: The init function to run before parameterized tests.
+ * @param_exit: The exit function to run after parameterized tests.
  *
  * A test case is a function with the signature,
  * ``void (*)(struct kunit *)``
@@ -129,6 +131,13 @@ struct kunit_case {
 	const void* (*generate_params)(const void *prev, char *desc);
 	struct kunit_attributes attr;
 
+	/*
+	 * Optional user-defined functions: one to register shared resources once
+	 * before the parameterized test series, and another to release them after.
+	 */
+	int (*param_init)(struct kunit *test);
+	void (*param_exit)(struct kunit *test);
+
 	/* private: internal use only. */
 	enum kunit_status status;
 	char *module_name;
@@ -218,6 +227,27 @@ static inline char *kunit_status_to_ok_not_ok(enum kunit_status status)
 		  .generate_params = gen_params,				\
 		  .attr = attributes, .module_name = KBUILD_MODNAME}
 
+/**
+ * KUNIT_CASE_PARAM_WITH_INIT() - Define a parameterized KUnit test case with custom
+ * init and exit functions.
+ * @test_name: The function implementing the test case.
+ * @gen_params: The function to generate parameters for the test case.
+ * @init: The init function to run before parameterized tests.
+ * @exit: The exit function to run after parameterized tests.
+ *
+ * Provides the option to register init and exit functions that take in the
+ * parent of the parameterized tests and run once before and once after the
+ * parameterized test series. The init function can be used to add any resources
+ * to share between the parameterized tests or to pass parameter arrays. The
+ * exit function can be used to clean up any resources that are not managed by
+ * the test.
+ */
+#define KUNIT_CASE_PARAM_WITH_INIT(test_name, gen_params, init, exit)		\
+		{ .run_case = test_name, .name = #test_name,			\
+		  .generate_params = gen_params,				\
+		  .param_init = init, .param_exit = exit,			\
+		  .module_name = KBUILD_MODNAME}
+
 /**
  * struct kunit_suite - describes a related collection of &struct kunit_case
  *
@@ -269,7 +299,8 @@ struct kunit_suite_set {
  * @priv: for user to store arbitrary data. Commonly used to pass data
  *	  created in the init function (see &struct kunit_suite).
  * @parent: for user to store data that they want to shared across
- *	    parameterized tests.
+ *	    parameterized tests. Typically, the data is provided in
+ *	    the param_init function (see &struct kunit_case).
  *
  * Used to store information about the current context under which the test
  * is running. Most of this data is private and should only be accessed
diff --git a/lib/kunit/test.c b/lib/kunit/test.c
index 4d6a39eb2c80..d80b5990d85d 100644
--- a/lib/kunit/test.c
+++ b/lib/kunit/test.c
@@ -641,6 +641,19 @@ static void kunit_accumulate_stats(struct kunit_result_stats *total,
 	total->total += add.total;
 }
 
+static void __kunit_init_parent_test(struct kunit_case *test_case, struct kunit *test)
+{
+	if (test_case->param_init) {
+		int err = test_case->param_init(test);
+
+		if (err) {
+			kunit_err(test_case, KUNIT_SUBTEST_INDENT KUNIT_SUBTEST_INDENT
+				"# failed to initialize parent parameter test.");
+			test_case->status = KUNIT_FAILURE;
+		}
+	}
+}
+
 int kunit_run_tests(struct kunit_suite *suite)
 {
 	char param_desc[KUNIT_PARAM_DESC_SIZE];
@@ -668,6 +681,8 @@ int kunit_run_tests(struct kunit_suite *suite)
 		struct kunit_result_stats param_stats = { 0 };
 
 		kunit_init_test(&test, test_case->name, test_case->log);
+		__kunit_init_parent_test(test_case, &test);
+
 		if (test_case->status == KUNIT_SKIPPED) {
 			/* Test marked as skip */
 			test.status = KUNIT_SKIPPED;
@@ -677,7 +692,7 @@ int kunit_run_tests(struct kunit_suite *suite)
 			test_case->status = KUNIT_SKIPPED;
 			kunit_run_case_catch_errors(suite, test_case, &test);
 			kunit_update_stats(&param_stats, test.status);
-		} else {
+		} else if (test_case->status != KUNIT_FAILURE) {
 			/* Get initial param. */
 			param_desc[0] = '\0';
 			/* TODO: Make generate_params try-catch */
@@ -727,6 +742,12 @@ int kunit_run_tests(struct kunit_suite *suite)
 
 		kunit_update_stats(&suite_stats, test_case->status);
 		kunit_accumulate_stats(&total_stats, param_stats);
+		/*
+		 * TODO: Put into a try catch. Since we don't need suite->exit
+		 * for it we can't reuse kunit_try_run_cleanup for this yet.
+		 */
+		if (test_case->param_exit)
+			test_case->param_exit(&test);
 		/* TODO: Put this kunit_cleanup into a try-catch. */
 		kunit_cleanup(&test);
 	}
-- 
2.50.1.552.g942d659e1b-goog