macros/
kunit.rs

1// SPDX-License-Identifier: GPL-2.0
2
3//! Procedural macro to run KUnit tests using a user-space like syntax.
4//!
5//! Copyright (c) 2023 José Expósito <jose.exposito89@gmail.com>
6
7use proc_macro::{Delimiter, Group, TokenStream, TokenTree};
8use std::fmt::Write;
9
10pub(crate) fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
11    let attr = attr.to_string();
12
13    if attr.is_empty() {
14        panic!("Missing test name in `#[kunit_tests(test_name)]` macro")
15    }
16
17    if attr.len() > 255 {
18        panic!("The test suite name `{attr}` exceeds the maximum length of 255 bytes")
19    }
20
21    let mut tokens: Vec<_> = ts.into_iter().collect();
22
23    // Scan for the `mod` keyword.
24    tokens
25        .iter()
26        .find_map(|token| match token {
27            TokenTree::Ident(ident) => match ident.to_string().as_str() {
28                "mod" => Some(true),
29                _ => None,
30            },
31            _ => None,
32        })
33        .expect("`#[kunit_tests(test_name)]` attribute should only be applied to modules");
34
35    // Retrieve the main body. The main body should be the last token tree.
36    let body = match tokens.pop() {
37        Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Brace => group,
38        _ => panic!("Cannot locate main body of module"),
39    };
40
41    // Get the functions set as tests. Search for `[test]` -> `fn`.
42    let mut body_it = body.stream().into_iter();
43    let mut tests = Vec::new();
44    while let Some(token) = body_it.next() {
45        match token {
46            TokenTree::Group(ident) if ident.to_string() == "[test]" => match body_it.next() {
47                Some(TokenTree::Ident(ident)) if ident.to_string() == "fn" => {
48                    let test_name = match body_it.next() {
49                        Some(TokenTree::Ident(ident)) => ident.to_string(),
50                        _ => continue,
51                    };
52                    tests.push(test_name);
53                }
54                _ => continue,
55            },
56            _ => (),
57        }
58    }
59
60    // Add `#[cfg(CONFIG_KUNIT="y")]` before the module declaration.
61    let config_kunit = "#[cfg(CONFIG_KUNIT=\"y\")]".to_owned().parse().unwrap();
62    tokens.insert(
63        0,
64        TokenTree::Group(Group::new(Delimiter::None, config_kunit)),
65    );
66
67    // Generate the test KUnit test suite and a test case for each `#[test]`.
68    // The code generated for the following test module:
69    //
70    // ```
71    // #[kunit_tests(kunit_test_suit_name)]
72    // mod tests {
73    //     #[test]
74    //     fn foo() {
75    //         assert_eq!(1, 1);
76    //     }
77    //
78    //     #[test]
79    //     fn bar() {
80    //         assert_eq!(2, 2);
81    //     }
82    // }
83    // ```
84    //
85    // Looks like:
86    //
87    // ```
88    // unsafe extern "C" fn kunit_rust_wrapper_foo(_test: *mut ::kernel::bindings::kunit) { foo(); }
89    // unsafe extern "C" fn kunit_rust_wrapper_bar(_test: *mut ::kernel::bindings::kunit) { bar(); }
90    //
91    // static mut TEST_CASES: [::kernel::bindings::kunit_case; 3] = [
92    //     ::kernel::kunit::kunit_case(::kernel::c_str!("foo"), kunit_rust_wrapper_foo),
93    //     ::kernel::kunit::kunit_case(::kernel::c_str!("bar"), kunit_rust_wrapper_bar),
94    //     ::kernel::kunit::kunit_case_null(),
95    // ];
96    //
97    // ::kernel::kunit_unsafe_test_suite!(kunit_test_suit_name, TEST_CASES);
98    // ```
99    let mut kunit_macros = "".to_owned();
100    let mut test_cases = "".to_owned();
101    let mut assert_macros = "".to_owned();
102    let path = crate::helpers::file();
103    for test in &tests {
104        let kunit_wrapper_fn_name = format!("kunit_rust_wrapper_{test}");
105        // An extra `use` is used here to reduce the length of the message.
106        let kunit_wrapper = format!(
107            "unsafe extern \"C\" fn {kunit_wrapper_fn_name}(_test: *mut ::kernel::bindings::kunit) {{ use ::kernel::kunit::is_test_result_ok; assert!(is_test_result_ok({test}())); }}",
108        );
109        writeln!(kunit_macros, "{kunit_wrapper}").unwrap();
110        writeln!(
111            test_cases,
112            "    ::kernel::kunit::kunit_case(::kernel::c_str!(\"{test}\"), {kunit_wrapper_fn_name}),"
113        )
114        .unwrap();
115        writeln!(
116            assert_macros,
117            r#"
118/// Overrides the usual [`assert!`] macro with one that calls KUnit instead.
119#[allow(unused)]
120macro_rules! assert {{
121    ($cond:expr $(,)?) => {{{{
122        kernel::kunit_assert!("{test}", "{path}", 0, $cond);
123    }}}}
124}}
125
126/// Overrides the usual [`assert_eq!`] macro with one that calls KUnit instead.
127#[allow(unused)]
128macro_rules! assert_eq {{
129    ($left:expr, $right:expr $(,)?) => {{{{
130        kernel::kunit_assert_eq!("{test}", "{path}", 0, $left, $right);
131    }}}}
132}}
133        "#
134        )
135        .unwrap();
136    }
137
138    writeln!(kunit_macros).unwrap();
139    writeln!(
140        kunit_macros,
141        "static mut TEST_CASES: [::kernel::bindings::kunit_case; {}] = [\n{test_cases}    ::kernel::kunit::kunit_case_null(),\n];",
142        tests.len() + 1
143    )
144    .unwrap();
145
146    writeln!(
147        kunit_macros,
148        "::kernel::kunit_unsafe_test_suite!({attr}, TEST_CASES);"
149    )
150    .unwrap();
151
152    // Remove the `#[test]` macros.
153    // We do this at a token level, in order to preserve span information.
154    let mut new_body = vec![];
155    let mut body_it = body.stream().into_iter();
156
157    while let Some(token) = body_it.next() {
158        match token {
159            TokenTree::Punct(ref c) if c.as_char() == '#' => match body_it.next() {
160                Some(TokenTree::Group(group)) if group.to_string() == "[test]" => (),
161                Some(next) => {
162                    new_body.extend([token, next]);
163                }
164                _ => {
165                    new_body.push(token);
166                }
167            },
168            _ => {
169                new_body.push(token);
170            }
171        }
172    }
173
174    let mut final_body = TokenStream::new();
175    final_body.extend::<TokenStream>(assert_macros.parse().unwrap());
176    final_body.extend(new_body);
177    final_body.extend::<TokenStream>(kunit_macros.parse().unwrap());
178
179    tokens.push(TokenTree::Group(Group::new(Delimiter::Brace, final_body)));
180
181    tokens.into_iter().collect()
182}
macros/kunit.rs

macros/
kunit.rs
