[PATCH 2/3] kunit: add ability to register a simple cleanup function

[benjamin@sipsolutions.net]

From: Benjamin Berg <benjamin.berg@intel.com>

This is useful to e.g. automatically free resources at the end of the
test, without having to deal with kunit resource objects directly.

The whole point of doing this is that the callback is guaranteed to
happen in case the test aborts (due to an assertion). As such, use
assertions internally rather than requiring further error checking by
the API user.

Signed-off-by: Benjamin Berg <benjamin.berg@intel.com>
---
 include/kunit/test.h | 20 ++++++++++++++++++++
 lib/kunit/test.c     | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/include/kunit/test.h b/include/kunit/test.h
index 519b90261c72..ab1dacf1c9f4 100644
--- a/include/kunit/test.h
+++ b/include/kunit/test.h
@@ -373,6 +373,26 @@ static inline void *kunit_kcalloc(struct kunit *test, size_t n, size_t size, gfp
 	return kunit_kmalloc_array(test, n, size, gfp | __GFP_ZERO);
 }
 
+typedef void (*kunit_cleanup_t)(const void *);
+
+/**
+ * kunit_add_cleanup() - Add post-test cleanup action.
+ * @test: The test case to which the resource belongs.
+ * @cleanup_func: function to call at end of test.
+ * @data: data to pass to @free_func.
+ * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
+ *
+ * This adds a cleanup action to be executed after the test completes.
+ * Internally this is handled using a *test managed resource*.
+ *
+ * This function will abort the test on failure.
+ *
+ * Note: KUnit needs to allocate memory for a kunit_resource object. You must
+ * specify an @internal_gfp that is compatible with the current context.
+ */
+void kunit_add_cleanup(struct kunit *test, kunit_cleanup_t cleanup_func,
+		       const void *data, gfp_t internal_gfp);
+
 void kunit_cleanup(struct kunit *test);
 
 void __printf(2, 3) kunit_log_append(char *log, const char *fmt, ...);
diff --git a/lib/kunit/test.c b/lib/kunit/test.c
index c9e15bb60058..72d956dfc324 100644
--- a/lib/kunit/test.c
+++ b/lib/kunit/test.c
@@ -719,6 +719,43 @@ void *kunit_kmalloc_array(struct kunit *test, size_t n, size_t size, gfp_t gfp)
 }
 EXPORT_SYMBOL_GPL(kunit_kmalloc_array);
 
+struct kunit_auto_cleanup {
+	struct kunit_resource resource;
+	kunit_cleanup_t cleanup_func;
+};
+
+static void kunit_auto_cleanup_free(struct kunit_resource *res)
+{
+	struct kunit_auto_cleanup *cleanup;
+
+	cleanup = container_of(res, struct kunit_auto_cleanup, resource);
+
+	cleanup->cleanup_func(cleanup->resource.data);
+}
+
+void kunit_add_cleanup(struct kunit *test, kunit_cleanup_t cleanup_func,
+		       const void *data, gfp_t internal_gfp)
+{
+	struct kunit_auto_cleanup *res;
+
+	KUNIT_ASSERT_NOT_NULL_MSG(test, cleanup_func,
+				  "Cleanup function must not be NULL");
+
+	res = kzalloc(sizeof(*res), internal_gfp);
+	if (!res) {
+		cleanup_func(data);
+		KUNIT_ASSERT_FAILURE(test, "Could not allocate resource for cleanup");
+	}
+
+	res->cleanup_func = cleanup_func;
+	res->resource.should_kfree = true;
+
+	/* Cannot fail as init is NULL */
+	__kunit_add_resource(test, NULL, kunit_auto_cleanup_free,
+			     &res->resource, (void *)data);
+}
+EXPORT_SYMBOL_GPL(kunit_add_cleanup);
+
 static inline bool kunit_kfree_match(struct kunit *test,
 				     struct kunit_resource *res, void *match_data)
 {
-- 
2.39.2