Re: [PATCH 2/3] rust: macros: add macro to easily run KUnit tests

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

On Thu, Jul 20, 2023 at 02:38:53PM +0800, David Gow wrote:
> From: José Expósito <jose.exposito89@gmail.com>
> 
> Add a new procedural macro (`#[kunit_tests(kunit_test_suit_name)]`) to
> run KUnit tests using a user-space like syntax.
> 
> The macro, that should be used on modules, transforms every `#[test]`
> in a `kunit_case!` and adds a `kunit_unsafe_test_suite!` registering
> all of them.
> 
> The only difference with user-space tests is that instead of using
> `#[cfg(test)]`, `#[kunit_tests(kunit_test_suit_name)]` is used.
> 
> Note that `#[cfg(CONFIG_KUNIT)]` is added so the test module is not
> compiled when `CONFIG_KUNIT` is set to `n`.
> 
> Reviewed-by: David Gow <davidgow@google.com>
> Signed-off-by: José Expósito <jose.exposito89@gmail.com>
> ---
>  MAINTAINERS          |   1 +
>  rust/kernel/kunit.rs |  11 ++++
>  rust/macros/kunit.rs | 149 +++++++++++++++++++++++++++++++++++++++++++++++++++
>  rust/macros/lib.rs   |  29 ++++++++++
>  4 files changed, 190 insertions(+)
> 
> diff --git a/MAINTAINERS b/MAINTAINERS
> index 2a942fe59144..c32ba6b29a96 100644
> --- a/MAINTAINERS
> +++ b/MAINTAINERS
> @@ -11373,6 +11373,7 @@ F:	Documentation/dev-tools/kunit/
>  F:	include/kunit/
>  F:	lib/kunit/
>  F:	rust/kernel/kunit.rs
> +F:	rust/macros/kunit.rs
>  F:	scripts/rustdoc_test_*
>  F:	tools/testing/kunit/
>  
> diff --git a/rust/kernel/kunit.rs b/rust/kernel/kunit.rs
> index 4cffc71e463b..44ea67028316 100644
> --- a/rust/kernel/kunit.rs
> +++ b/rust/kernel/kunit.rs
> @@ -40,6 +40,8 @@ pub fn info(args: fmt::Arguments<'_>) {
>      }
>  }
>  
> +use macros::kunit_tests;
> +
>  /// Asserts that a boolean expression is `true` at runtime.
>  ///
>  /// Public but hidden since it should only be used from generated tests.
> @@ -253,3 +255,12 @@ macro_rules! kunit_unsafe_test_suite {
>          };
>      };
>  }
> +
> +#[kunit_tests(rust_kernel_kunit)]
> +mod tests {
> +    #[test]
> +    fn rust_test_kunit_kunit_tests() {
> +        let running = true;
> +        assert_eq!(running, true);
> +    }
> +}
> diff --git a/rust/macros/kunit.rs b/rust/macros/kunit.rs
> new file mode 100644
> index 000000000000..69dac253232f
> --- /dev/null
> +++ b/rust/macros/kunit.rs
> @@ -0,0 +1,149 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +//! Procedural macro to run KUnit tests using a user-space like syntax.
> +//!
> +//! Copyright (c) 2023 José Expósito <jose.exposito89@gmail.com>
> +
> +use proc_macro::{Delimiter, Group, TokenStream, TokenTree};
> +use std::fmt::Write;
> +
> +pub(crate) fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
> +    if attr.to_string().is_empty() {
> +        panic!("Missing test name in #[kunit_tests(test_name)] macro")
> +    }
> +
> +    let mut tokens: Vec<_> = ts.into_iter().collect();
> +
> +    // Scan for the "mod" keyword.
> +    tokens
> +        .iter()
> +        .find_map(|token| match token {
> +            TokenTree::Ident(ident) => match ident.to_string().as_str() {
> +                "mod" => Some(true),
> +                _ => None,
> +            },
> +            _ => None,
> +        })
> +        .expect("#[kunit_tests(test_name)] attribute should only be applied to modules");
> +
> +    // Retrieve the main body. The main body should be the last token tree.
> +    let body = match tokens.pop() {
> +        Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Brace => group,
> +        _ => panic!("cannot locate main body of module"),
> +    };
> +
> +    // Get the functions set as tests. Search for `[test]` -> `fn`.
> +    let mut body_it = body.stream().into_iter();
> +    let mut tests = Vec::new();
> +    while let Some(token) = body_it.next() {
> +        match token {
> +            TokenTree::Group(ident) if ident.to_string() == "[test]" => match body_it.next() {
> +                Some(TokenTree::Ident(ident)) if ident.to_string() == "fn" => {
> +                    let test_name = match body_it.next() {
> +                        Some(TokenTree::Ident(ident)) => ident.to_string(),
> +                        _ => continue,
> +                    };
> +                    tests.push(test_name);
> +                }
> +                _ => continue,
> +            },
> +            _ => (),
> +        }
> +    }
> +
> +    // Add `#[cfg(CONFIG_KUNIT)]` before the module declaration.
> +    let config_kunit = "#[cfg(CONFIG_KUNIT)]".to_owned().parse().unwrap();
> +    tokens.insert(
> +        0,
> +        TokenTree::Group(Group::new(Delimiter::None, config_kunit)),
> +    );
> +
> +    // Generate the test KUnit test suite and a test case for each `#[test]`.
> +    // The code generated for the following test module:
> +    //
> +    // ```
> +    // #[kunit_tests(kunit_test_suit_name)]
> +    // mod tests {
> +    //     #[test]
> +    //     fn foo() {
> +    //         assert_eq!(1, 1);
> +    //     }
> +    //
> +    //     #[test]
> +    //     fn bar() {
> +    //         assert_eq!(2, 2);
> +    //     }
> +    // ```
> +    //
> +    // Looks like:
> +    //
> +    // ```
> +    // unsafe extern "C" fn kunit_rust_wrapper_foo(_test: *mut kernel::bindings::kunit) {
> +    //     foo();
> +    // }
> +    // static mut KUNIT_CASE_FOO: kernel::bindings::kunit_case =
> +    //     kernel::kunit_case!(foo, kunit_rust_wrapper_foo);
> +    //
> +    // unsafe extern "C" fn kunit_rust_wrapper_bar(_test: * mut kernel::bindings::kunit) {
> +    //     bar();
> +    // }
> +    // static mut KUNIT_CASE_BAR: kernel::bindings::kunit_case =
> +    //     kernel::kunit_case!(bar, kunit_rust_wrapper_bar);
> +    //
> +    // static mut KUNIT_CASE_NULL: kernel::bindings::kunit_case = kernel::kunit_case!();
> +    //
> +    // static mut TEST_CASES : &mut[kernel::bindings::kunit_case] = unsafe {
> +    //     &mut [KUNIT_CASE_FOO, KUNIT_CASE_BAR, KUNIT_CASE_NULL]
> +    // };
> +    //
> +    // kernel::kunit_unsafe_test_suite!(kunit_test_suit_name, TEST_CASES);
> +    // ```
> +    let mut kunit_macros = "".to_owned();
> +    let mut test_cases = "".to_owned();
> +    for test in tests {
> +        let kunit_wrapper_fn_name = format!("kunit_rust_wrapper_{}", test);
> +        let kunit_case_name = format!("KUNIT_CASE_{}", test.to_uppercase());
> +        let kunit_wrapper = format!(
> +            "unsafe extern \"C\" fn {}(_test: *mut kernel::bindings::kunit) {{ {}(); }}",
> +            kunit_wrapper_fn_name, test
> +        );
> +        let kunit_case = format!(
> +            "static mut {}: kernel::bindings::kunit_case = kernel::kunit_case!({}, {});",
> +            kunit_case_name, test, kunit_wrapper_fn_name
> +        );
> +        writeln!(kunit_macros, "{kunit_wrapper}").unwrap();
> +        writeln!(kunit_macros, "{kunit_case}").unwrap();
> +        writeln!(test_cases, "{kunit_case_name},").unwrap();
> +    }
> +
> +    writeln!(
> +        kunit_macros,
> +        "static mut KUNIT_CASE_NULL: kernel::bindings::kunit_case = kernel::kunit_case!();"
> +    )
> +    .unwrap();
> +
> +    writeln!(
> +        kunit_macros,
> +        "static mut TEST_CASES : &mut[kernel::bindings::kunit_case] = unsafe {{ &mut[{test_cases} KUNIT_CASE_NULL] }};"
> +    )
> +    .unwrap();
> +
> +    writeln!(
> +        kunit_macros,
> +        "kernel::kunit_unsafe_test_suite!({attr}, TEST_CASES);"
> +    )
> +    .unwrap();
> +
> +    let new_body: TokenStream = vec![body.stream(), kunit_macros.parse().unwrap()]
> +        .into_iter()
> +        .collect();
> +
> +    // Remove the `#[test]` macros.
> +    let new_body = new_body.to_string().replace("#[test]", "");

I've played this with some extra tests, one thing I notice is that
Span/code location information is lost if we do this, for example, if I
have a compile error in the test code (I introduced one on purpose in
the `rust_kernel_kunit` test), I will get information like:

	error[E0384]: cannot assign twice to immutable variable `running`
	   --> ../rust/kernel/kunit.rs:329:1
	    |
	329 | #[kunit_tests(rust_kernel_kunit)]
	    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
	    | |
	    | cannot assign twice to immutable variable
	    | help: consider making this binding mutable: `mut running`

the location information is not very usefull. However if we do the
following:

diff --git a/rust/macros/kunit.rs b/rust/macros/kunit.rs
index 69dac253232f..913879765d24 100644
--- a/rust/macros/kunit.rs
+++ b/rust/macros/kunit.rs
@@ -134,15 +134,29 @@ pub(crate) fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
     )
     .unwrap();
 
-    let new_body: TokenStream = vec![body.stream(), kunit_macros.parse().unwrap()]
-        .into_iter()
-        .collect();
-
     // Remove the `#[test]` macros.
-    let new_body = new_body.to_string().replace("#[test]", "");
+    let mut new_body = vec![];
+    let mut body_it = body.stream().into_iter();
+
+    while let Some(token) = body_it.next() {
+        match token {
+            TokenTree::Punct(ref c) if c.as_char() == '#' => {
+                match body_it.next() {
+                    Some(TokenTree::Group(g)) if g.to_string() == "[test]" => (),
+                    Some(next) => { new_body.extend([token, next]);},
+                    _ => {new_body.push(token);},
+                }
+            }
+            _ => { new_body.push(token); }
+        }
+    }
+
+    let mut new_body = TokenStream::from_iter(new_body);
+    new_body.extend::<TokenStream>(kunit_macros.parse().unwrap());
+
     tokens.push(TokenTree::Group(Group::new(
         Delimiter::Brace,
-        new_body.parse().unwrap(),
+        new_body
     )));
 
     tokens.into_iter().collect()


then we get better information:

	   --> ../rust/kernel/kunit.rs:335:13
	    |
	335 |         let running = true;
	    |             ^^^^^^^
	    |
	    = help: maybe it is overwritten before being read?
	    = note: `#[warn(unused_assignments)]` on by default

	error[E0384]: cannot assign twice to immutable variable `running`
	   --> ../rust/kernel/kunit.rs:336:9
	    |
	335 |         let running = true;
	    |             -------
	    |             |
	    |             first assignment to `running`
	    |             help: consider making this binding mutable: `mut running`
	336 |         running = false;
	    |         ^^^^^^^^^^^^^^^ cannot assign twice to immutable variable


Regards,
Boqun

> +    tokens.push(TokenTree::Group(Group::new(
> +        Delimiter::Brace,
> +        new_body.parse().unwrap(),
> +    )));
> +
> +    tokens.into_iter().collect()
> +}
> diff --git a/rust/macros/lib.rs b/rust/macros/lib.rs
> index 3fc74cb4ea19..cd0b720514ff 100644
> --- a/rust/macros/lib.rs
> +++ b/rust/macros/lib.rs
> @@ -6,6 +6,7 @@
>  mod quote;
>  mod concat_idents;
>  mod helpers;
> +mod kunit;
>  mod module;
>  mod pin_data;
>  mod pinned_drop;
> @@ -246,3 +247,31 @@ pub fn pin_data(inner: TokenStream, item: TokenStream) -> TokenStream {
>  pub fn pinned_drop(args: TokenStream, input: TokenStream) -> TokenStream {
>      pinned_drop::pinned_drop(args, input)
>  }
> +
> +/// Registers a KUnit test suite and its test cases using a user-space like syntax.
> +///
> +/// This macro should be used on modules. If `CONFIG_KUNIT` (in `.config`) is `n`, the target module
> +/// is ignored.
> +///
> +/// # Examples
> +///
> +/// ```ignore
> +/// # use macros::kunit_tests;
> +///
> +/// #[kunit_tests(kunit_test_suit_name)]
> +/// mod tests {
> +///     #[test]
> +///     fn foo() {
> +///         assert_eq!(1, 1);
> +///     }
> +///
> +///     #[test]
> +///     fn bar() {
> +///         assert_eq!(2, 2);
> +///     }
> +/// }
> +/// ```
> +#[proc_macro_attribute]
> +pub fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
> +    kunit::kunit_tests(attr, ts)
> +}
> 
> -- 
> 2.41.0.255.g8b1d071c50-goog
>