Re: [PATCH 3/3] rust: kunit: allow to know if we are in a test

[Boqun Feng <boqun.feng@gmail.com>]

On Thu, Jul 20, 2023 at 02:38:54PM +0800, David Gow wrote:
> From: José Expósito <jose.exposito89@gmail.com>
> 
> In some cases, you need to call test-only code from outside the test
> case, for example, to mock a function or a module.
> 
> In order to check whether we are in a test or not, we need to test if
> `CONFIG_KUNIT` is set.
> Unfortunately, we cannot rely only on this condition because some
> distros compile KUnit in production kernels, so checking at runtime
> that `current->kunit_test != NULL` is required.
> 
> Note that the C function `kunit_get_current_test()` can not be used
> because it is not present in the current Rust tree yet. Once it is
> available we might want to change our Rust wrapper to use it.
> 
> This patch adds a function to know whether we are in a KUnit test or
> not and examples showing how to mock a function and a module.
> 
> Reviewed-by: David Gow <davidgow@google.com>
> Signed-off-by: José Expósito <jose.exposito89@gmail.com>
> ---
>  rust/kernel/kunit.rs | 78 ++++++++++++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 78 insertions(+)
> 
> diff --git a/rust/kernel/kunit.rs b/rust/kernel/kunit.rs
> index 44ea67028316..dcaac19bb108 100644
> --- a/rust/kernel/kunit.rs
> +++ b/rust/kernel/kunit.rs
> @@ -40,6 +40,8 @@ pub fn info(args: fmt::Arguments<'_>) {
>      }
>  }
>  
> +use crate::task::Task;
> +use core::ops::Deref;
>  use macros::kunit_tests;
>  
>  /// Asserts that a boolean expression is `true` at runtime.
> @@ -256,11 +258,87 @@ macro_rules! kunit_unsafe_test_suite {
>      };
>  }
>  
> +/// In some cases, you need to call test-only code from outside the test case, for example, to
> +/// create a function mock. This function can be invoked to know whether we are currently running a
> +/// KUnit test or not.
> +///
> +/// # Examples
> +///
> +/// This example shows how a function can be mocked to return a well-known value while testing:
> +///
> +/// ```
> +/// # use kernel::kunit::in_kunit_test;
> +/// #
> +/// fn fn_mock_example(n: i32) -> i32 {
> +///     if in_kunit_test() {
> +///         100
> +///     } else {
> +///         n + 1
> +///     }
> +/// }
> +///
> +/// let mock_res = fn_mock_example(5);
> +/// assert_eq!(mock_res, 100);
> +/// ```
> +///
> +/// Sometimes, you don't control the code that needs to be mocked. This example shows how the
> +/// `bindings` module can be mocked:
> +///
> +/// ```
> +/// // Import our mock naming it as the real module.
> +/// #[cfg(CONFIG_KUNIT)]
> +/// use bindings_mock_example as bindings;
> +///
> +/// // This module mocks `bindings`.
> +/// mod bindings_mock_example {
> +///     use kernel::kunit::in_kunit_test;
> +///     use kernel::bindings::u64_;
> +///
> +///     // Make the other binding functions available.
> +///     pub(crate) use kernel::bindings::*;
> +///
> +///     // Mock `ktime_get_boot_fast_ns` to return a well-known value when running a KUnit test.
> +///     pub(crate) unsafe fn ktime_get_boot_fast_ns() -> u64_ {
> +///         if in_kunit_test() {
> +///             1234
> +///         } else {
> +///             unsafe { kernel::bindings::ktime_get_boot_fast_ns() }
> +///         }
> +///     }
> +/// }
> +///
> +/// // This is the function we want to test. Since `bindings` has been mocked, we can use its
> +/// // functions seamlessly.
> +/// fn get_boot_ns() -> u64 {
> +///     unsafe { bindings::ktime_get_boot_fast_ns() }
> +/// }
> +///
> +/// let time = get_boot_ns();
> +/// assert_eq!(time, 1234);
> +/// ```
> +pub fn in_kunit_test() -> bool {
> +    if cfg!(CONFIG_KUNIT) {
> +        // SAFETY: By the type invariant, we know that `*Task::current().deref().0` is valid.
> +        let test = unsafe { (*Task::current().deref().0.get()).kunit_test };

Note here are two unsafe operations: `Task::current()` and the pointer
dereference. You can use the `current!()` macro here to avoid the first
unsafe operation here. Besides I think it'll be better if
in_kunit_test() is a safe method for `Task`? That will be easier for us
to track the usage of task_struct fields in Rust side. But I'm OK with
either way.

Regards,
Boqun

> +        !test.is_null()
> +    } else {
> +        false
> +    }
> +}
> +
>  #[kunit_tests(rust_kernel_kunit)]
>  mod tests {
> +    use super::*;
> +
>      #[test]
>      fn rust_test_kunit_kunit_tests() {
>          let running = true;
>          assert_eq!(running, true);
>      }
> +
> +    #[test]
> +    fn rust_test_kunit_in_kunit_test() {
> +        let in_kunit = in_kunit_test();
> +        assert_eq!(in_kunit, true);
> +    }
>  }
> 
> -- 
> 2.41.0.255.g8b1d071c50-goog
>