[PATCH 07/12] drm/xe/tests: Add helpers to call stubs out of KUnit context

[Michal Wajdeczko <michal.wajdeczko@intel.com>]

The KUNIT_STATIC_STUB_REDIRECT() allows to redirect function call but will
work only if the caller is in a KUnit context. To allow implementation of
the more complex test cases, add helpers that allow calling stubs also out
of a KUnit context.

Signed-off-by: Michal Wajdeczko <michal.wajdeczko@intel.com>
Cc: Lucas De Marchi <lucas.demarchi@intel.com>
---
 drivers/gpu/drm/xe/tests/xe_test.h | 74 ++++++++++++++++++++++++++++++
 1 file changed, 74 insertions(+)

diff --git a/drivers/gpu/drm/xe/tests/xe_test.h b/drivers/gpu/drm/xe/tests/xe_test.h
index b8fa409ce2b1..ffa8fa6c96f0 100644
--- a/drivers/gpu/drm/xe/tests/xe_test.h
+++ b/drivers/gpu/drm/xe/tests/xe_test.h
@@ -54,11 +54,85 @@ xe_cur_kunit_priv(enum xe_test_priv_id id)
 	return priv->id == id ? priv : NULL;
 }
 
+/**
+ * XE_TEST_REDIRECT() - Call a function stub if one exists.
+ * @stub: The pointer to the function stub
+ * @args: All of the arguments passed to this stub
+ *
+ * This is a function prologue which is used to allow calls to the current
+ * function to be redirected if a KUnit is running. If the stub is NULL or
+ * the KUnit is not running the function will continue execution as normal.
+ *
+ * Unlikely the KUNIT_STATIC_STUB_REDIRECT(), this redirection will work
+ * even if the caller is not in a KUnit context (like a worker thread).
+ *
+ * Example:
+ *
+ * .. code-block:: c
+ *
+ *	int (*stub)(int n);
+ *
+ *	int real_func(int n)
+ *	{
+ *		XE_TEST_REDIRECT(stub, n);
+ *		return n + 1;
+ *	}
+ *
+ *	int replacement_func(int n)
+ *	{
+ *		return n + 100;
+ *	}
+ *
+ *	void example_test(struct kunit *test)
+ *	{
+ *		stub = replacement_func;
+ *		KUNIT_EXPECT_EQ(test, real_func(1), 101);
+ *		KUNIT_EXPECT_EQ(test, real_func(1), 101);
+ *	}
+ */
+#define XE_TEST_REDIRECT(stub, args...)							\
+do {											\
+	typeof(stub) replacement = (stub);						\
+	if (XE_TEST_RUNNING()) {							\
+		if (unlikely(replacement)) {						\
+			pr_info(KUNIT_SUBTEST_INDENT "# %s: calling stub %ps\n",	\
+				__func__, replacement);					\
+			return replacement(args);					\
+		}									\
+	}										\
+} while (0)
+
+static inline void __nullify_pointer(void *data)
+{
+	void **ptr = data;
+
+	*ptr = NULL;
+}
+
+/**
+ * XE_TEST_ACTIVATE_STUB() - Setup a function stub.
+ * @test: Test case that wants to setup a function stub
+ * @stub: The function stub pointer
+ * @replacement: The replacement function
+ *
+ * This helper setups a function stub with the replacement function.
+ * It will also automatically restore stub to NULL at the test end.
+ */
+#define XE_TEST_ACTIVATE_STUB(test, stub, replacement)	({					\
+	typeof(test) __test = (test);								\
+	typeof(stub) *__ptr = &(stub);								\
+	typecheck_pointer(*__ptr);								\
+	*__ptr = (replacement);									\
+	kunit_info(__test, "activated stub %s with %ps\n", __stringify(stub), *__ptr);		\
+	kunit_add_action_or_reset(__test, __nullify_pointer, __ptr);				\
+})
+
 #else /* if IS_ENABLED(CONFIG_DRM_XE_KUNIT_TEST) */
 
 #define XE_TEST_DECLARE(x)
 #define XE_TEST_ONLY(x) 0
 #define XE_TEST_RUNNING() false
+#define XE_TEST_REDIRECT(...) do { } while (0)
 
 #define xe_cur_kunit_priv(_id) NULL
 
-- 
2.43.0