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::collections::HashMap;
9use std::fmt::Write;
10
11pub(crate) fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
12    let attr = attr.to_string();
13
14    if attr.is_empty() {
15        panic!("Missing test name in `#[kunit_tests(test_name)]` macro")
16    }
17
18    if attr.len() > 255 {
19        panic!("The test suite name `{attr}` exceeds the maximum length of 255 bytes")
20    }
21
22    let mut tokens: Vec<_> = ts.into_iter().collect();
23
24    // Scan for the `mod` keyword.
25    tokens
26        .iter()
27        .find_map(|token| match token {
28            TokenTree::Ident(ident) => match ident.to_string().as_str() {
29                "mod" => Some(true),
30                _ => None,
31            },
32            _ => None,
33        })
34        .expect("`#[kunit_tests(test_name)]` attribute should only be applied to modules");
35
36    // Retrieve the main body. The main body should be the last token tree.
37    let body = match tokens.pop() {
38        Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Brace => group,
39        _ => panic!("Cannot locate main body of module"),
40    };
41
42    // Get the functions set as tests. Search for `[test]` -> `fn`.
43    let mut body_it = body.stream().into_iter();
44    let mut tests = Vec::new();
45    let mut attributes: HashMap<String, TokenStream> = HashMap::new();
46    while let Some(token) = body_it.next() {
47        match token {
48            TokenTree::Punct(ref p) if p.as_char() == '#' => match body_it.next() {
49                Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Bracket => {
50                    if let Some(TokenTree::Ident(name)) = g.stream().into_iter().next() {
51                        // Collect attributes because we need to find which are tests. We also
52                        // need to copy `cfg` attributes so tests can be conditionally enabled.
53                        attributes
54                            .entry(name.to_string())
55                            .or_default()
56                            .extend([token, TokenTree::Group(g)]);
57                    }
58                    continue;
59                }
60                _ => (),
61            },
62            TokenTree::Ident(i) if i.to_string() == "fn" && attributes.contains_key("test") => {
63                if let Some(TokenTree::Ident(test_name)) = body_it.next() {
64                    tests.push((test_name, attributes.remove("cfg").unwrap_or_default()))
65                }
66            }
67
68            _ => (),
69        }
70        attributes.clear();
71    }
72
73    // Add `#[cfg(CONFIG_KUNIT="y")]` before the module declaration.
74    let config_kunit = "#[cfg(CONFIG_KUNIT=\"y\")]".to_owned().parse().unwrap();
75    tokens.insert(
76        0,
77        TokenTree::Group(Group::new(Delimiter::None, config_kunit)),
78    );
79
80    // Generate the test KUnit test suite and a test case for each `#[test]`.
81    // The code generated for the following test module:
82    //
83    // ```
84    // #[kunit_tests(kunit_test_suit_name)]
85    // mod tests {
86    //     #[test]
87    //     fn foo() {
88    //         assert_eq!(1, 1);
89    //     }
90    //
91    //     #[test]
92    //     fn bar() {
93    //         assert_eq!(2, 2);
94    //     }
95    // }
96    // ```
97    //
98    // Looks like:
99    //
100    // ```
101    // unsafe extern "C" fn kunit_rust_wrapper_foo(_test: *mut ::kernel::bindings::kunit) { foo(); }
102    // unsafe extern "C" fn kunit_rust_wrapper_bar(_test: *mut ::kernel::bindings::kunit) { bar(); }
103    //
104    // static mut TEST_CASES: [::kernel::bindings::kunit_case; 3] = [
105    //     ::kernel::kunit::kunit_case(::kernel::c_str!("foo"), kunit_rust_wrapper_foo),
106    //     ::kernel::kunit::kunit_case(::kernel::c_str!("bar"), kunit_rust_wrapper_bar),
107    //     ::kernel::kunit::kunit_case_null(),
108    // ];
109    //
110    // ::kernel::kunit_unsafe_test_suite!(kunit_test_suit_name, TEST_CASES);
111    // ```
112    let mut kunit_macros = "".to_owned();
113    let mut test_cases = "".to_owned();
114    let mut assert_macros = "".to_owned();
115    let path = crate::helpers::file();
116    let num_tests = tests.len();
117    for (test, cfg_attr) in tests {
118        let kunit_wrapper_fn_name = format!("kunit_rust_wrapper_{test}");
119        // Append any `cfg` attributes the user might have written on their tests so we don't
120        // attempt to call them when they are `cfg`'d out. An extra `use` is used here to reduce
121        // the length of the assert message.
122        let kunit_wrapper = format!(
123            r#"unsafe extern "C" fn {kunit_wrapper_fn_name}(_test: *mut ::kernel::bindings::kunit)
124            {{
125                (*_test).status = ::kernel::bindings::kunit_status_KUNIT_SKIPPED;
126                {cfg_attr} {{
127                    (*_test).status = ::kernel::bindings::kunit_status_KUNIT_SUCCESS;
128                    use ::kernel::kunit::is_test_result_ok;
129                    assert!(is_test_result_ok({test}()));
130                }}
131            }}"#,
132        );
133        writeln!(kunit_macros, "{kunit_wrapper}").unwrap();
134        writeln!(
135            test_cases,
136            "    ::kernel::kunit::kunit_case(::kernel::c_str!(\"{test}\"), {kunit_wrapper_fn_name}),"
137        )
138        .unwrap();
139        writeln!(
140            assert_macros,
141            r#"
142/// Overrides the usual [`assert!`] macro with one that calls KUnit instead.
143#[allow(unused)]
144macro_rules! assert {{
145    ($cond:expr $(,)?) => {{{{
146        kernel::kunit_assert!("{test}", "{path}", 0, $cond);
147    }}}}
148}}
149
150/// Overrides the usual [`assert_eq!`] macro with one that calls KUnit instead.
151#[allow(unused)]
152macro_rules! assert_eq {{
153    ($left:expr, $right:expr $(,)?) => {{{{
154        kernel::kunit_assert_eq!("{test}", "{path}", 0, $left, $right);
155    }}}}
156}}
157        "#
158        )
159        .unwrap();
160    }
161
162    writeln!(kunit_macros).unwrap();
163    writeln!(
164        kunit_macros,
165        "static mut TEST_CASES: [::kernel::bindings::kunit_case; {}] = [\n{test_cases}    ::kernel::kunit::kunit_case_null(),\n];",
166        num_tests + 1
167    )
168    .unwrap();
169
170    writeln!(
171        kunit_macros,
172        "::kernel::kunit_unsafe_test_suite!({attr}, TEST_CASES);"
173    )
174    .unwrap();
175
176    // Remove the `#[test]` macros.
177    // We do this at a token level, in order to preserve span information.
178    let mut new_body = vec![];
179    let mut body_it = body.stream().into_iter();
180
181    while let Some(token) = body_it.next() {
182        match token {
183            TokenTree::Punct(ref c) if c.as_char() == '#' => match body_it.next() {
184                Some(TokenTree::Group(group)) if group.to_string() == "[test]" => (),
185                Some(next) => {
186                    new_body.extend([token, next]);
187                }
188                _ => {
189                    new_body.push(token);
190                }
191            },
192            _ => {
193                new_body.push(token);
194            }
195        }
196    }
197
198    let mut final_body = TokenStream::new();
199    final_body.extend::<TokenStream>(assert_macros.parse().unwrap());
200    final_body.extend(new_body);
201    final_body.extend::<TokenStream>(kunit_macros.parse().unwrap());
202
203    tokens.push(TokenTree::Group(Group::new(Delimiter::Brace, final_body)));
204
205    tokens.into_iter().collect()
206}
macros/kunit.rs

macros/
kunit.rs
