[PATCH v2 3/6] lib: Add KUnit tests for union-find implementation

[Kuan-Wei Chiu <visitorckw@gmail.com>]

Introduce a KUnit test suite for the union-find data structure. The
tests verify the functionality and correctness of the union and find
operations, including edge cases such as handling duplicate unions.
The addition of KUnit tests enhances the robustness of the union-find
implementation by ensuring its correctness under various scenarios.

Signed-off-by: Kuan-Wei Chiu <visitorckw@gmail.com>
---
 MAINTAINERS            |  1 +
 lib/Kconfig.debug      | 12 +++++++
 lib/Makefile           |  1 +
 lib/union_find_kunit.c | 74 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 88 insertions(+)
 create mode 100644 lib/union_find_kunit.c

diff --git a/MAINTAINERS b/MAINTAINERS
index a097afd76ded..e503802c1549 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -23801,6 +23801,7 @@ F:	Documentation/core-api/union_find.rst
 F:	Documentation/translations/zh_CN/core-api/union_find.rst
 F:	include/linux/union_find.h
 F:	lib/union_find.c
+F:	lib/union_find_kunit.c
 
 UNIVERSAL FLASH STORAGE HOST CONTROLLER DRIVER
 R:	Alim Akhtar <alim.akhtar@samsung.com>
diff --git a/lib/Kconfig.debug b/lib/Kconfig.debug
index 7315f643817a..f8c09dd0590b 100644
--- a/lib/Kconfig.debug
+++ b/lib/Kconfig.debug
@@ -2841,6 +2841,18 @@ config SIPHASH_KUNIT_TEST
 	  This is intended to help people writing architecture-specific
 	  optimized versions.  If unsure, say N.
 
+config UNION_FIND_KUNIT_TEST
+	tristate "KUnit Test for Union find"
+	depends on KUNIT
+	default KUNIT_ALL_TESTS
+	help
+	  This option enables the KUnit tests for the union-find data structure.
+	  These tests verify the functionality and correctness of the union-find
+	  implementation, including union and find operations, as well as
+	  edge cases such as handling of duplicate unions.
+
+	  If unsure, say N
+
 config USERCOPY_KUNIT_TEST
 	tristate "KUnit Test for user/kernel boundary protections"
 	depends on KUNIT
diff --git a/lib/Makefile b/lib/Makefile
index 773adf88af41..03da92faf9b8 100644
--- a/lib/Makefile
+++ b/lib/Makefile
@@ -388,6 +388,7 @@ CFLAGS_fortify_kunit.o += $(call cc-disable-warning, stringop-truncation)
 CFLAGS_fortify_kunit.o += $(DISABLE_STRUCTLEAK_PLUGIN)
 obj-$(CONFIG_FORTIFY_KUNIT_TEST) += fortify_kunit.o
 obj-$(CONFIG_SIPHASH_KUNIT_TEST) += siphash_kunit.o
+obj-$(CONFIG_UNION_FIND_KUNIT_TEST) += union_find_kunit.o
 obj-$(CONFIG_USERCOPY_KUNIT_TEST) += usercopy_kunit.o
 
 obj-$(CONFIG_GENERIC_LIB_DEVMEM_IS_ALLOWED) += devmem_is_allowed.o
diff --git a/lib/union_find_kunit.c b/lib/union_find_kunit.c
new file mode 100644
index 000000000000..9bdf9e0e455e
--- /dev/null
+++ b/lib/union_find_kunit.c
@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: GPL-2.0-only
+
+#include <kunit/test.h>
+#include <linux/module.h>
+#include <linux/union_find.h>
+
+static void test_union_and_find(struct kunit *test)
+{
+	struct uf_node node1, node2, node3;
+	struct uf_node *root1, *root2, *root3;
+	bool merged;
+
+	/* Initialize the nodes */
+	uf_node_init(&node1);
+	uf_node_init(&node2);
+	uf_node_init(&node3);
+
+	/* Check the initial parent and rank */
+	KUNIT_ASSERT_PTR_EQ(test, uf_find(&node1), &node1);
+	KUNIT_ASSERT_PTR_EQ(test, uf_find(&node2), &node2);
+	KUNIT_ASSERT_PTR_EQ(test, uf_find(&node3), &node3);
+	KUNIT_ASSERT_EQ(test, node1.rank, 0);
+	KUNIT_ASSERT_EQ(test, node2.rank, 0);
+	KUNIT_ASSERT_EQ(test, node3.rank, 0);
+
+	/* Union node1 and node2 */
+	merged = uf_union(&node1, &node2);
+	KUNIT_ASSERT_TRUE(test, merged);
+
+	/* Assert that one of the nodes is now the parent of the other */
+	root1 = uf_find(&node1);
+	root2 = uf_find(&node2);
+	KUNIT_ASSERT_PTR_EQ(test, root1, root2);
+
+	/* Check rank after the first union */
+	if (root1 == &node1) {
+		KUNIT_ASSERT_EQ(test, node1.rank, 1);
+		KUNIT_ASSERT_EQ(test, node2.rank, 0);
+	} else {
+		KUNIT_ASSERT_EQ(test, node1.rank, 0);
+		KUNIT_ASSERT_EQ(test, node2.rank, 1);
+	}
+
+	/* Attempt to union node1 and node2 again and check for false return */
+	merged = uf_union(&node1, &node2);
+	KUNIT_ASSERT_FALSE(test, merged);
+
+	/* Union node3 with the result of the previous union (node1 and node2) */
+	uf_union(&node1, &node3);
+
+	/* Assert that all nodes have the same root */
+	root3 = uf_find(&node3);
+	KUNIT_ASSERT_PTR_EQ(test, root1, root3);
+
+	/* Check rank after the second union */
+	KUNIT_ASSERT_EQ(test, root1->rank, 1);
+	KUNIT_ASSERT_EQ(test, node3.rank, 0);
+}
+
+static struct kunit_case union_find_test_cases[] = {
+	KUNIT_CASE(test_union_and_find),
+	{}
+};
+
+static struct kunit_suite union_find_test_suite = {
+	.name = "union_find_test_suite",
+	.test_cases = union_find_test_cases,
+};
+
+kunit_test_suites(&union_find_test_suite);
+
+MODULE_AUTHOR("Kuan-Wei Chiu <visitorckw@gmail.com>");
+MODULE_DESCRIPTION("Union-find KUnit test suite");
+MODULE_LICENSE("GPL");
-- 
2.34.1