3 files changed, 52 insertions, 5 deletions
diff --git a/rust/macros/helpers.rs b/rust/macros/helpers.rs
index a3ee27e29a6f..e2602be402c1 100644
--- a/rust/macros/helpers.rs
+++ b/rust/macros/helpers.rs
@@ -86,3 +86,20 @@ pub(crate) fn function_name(input: TokenStream) -> Option<Ident> {
     }
     None
 }
+
+pub(crate) fn file() -> String {
+    #[cfg(not(CONFIG_RUSTC_HAS_SPAN_FILE))]
+    {
+        proc_macro::Span::call_site()
+            .source_file()
+            .path()
+            .to_string_lossy()
+            .into_owned()
+    }
+
+    #[cfg(CONFIG_RUSTC_HAS_SPAN_FILE)]
+    #[allow(clippy::incompatible_msrv)]
+    {
+        proc_macro::Span::call_site().file()
+    }
+}
diff --git a/rust/macros/kunit.rs b/rust/macros/kunit.rs
index b146c5894e92..971c2c3d849a 100644
--- a/rust/macros/kunit.rs
+++ b/rust/macros/kunit.rs
@@ -57,8 +57,8 @@ pub(crate) fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
         }
     }
 
-    // Add `#[cfg(CONFIG_KUNIT)]` before the module declaration.
-    let config_kunit = "#[cfg(CONFIG_KUNIT)]".to_owned().parse().unwrap();
+    // Add `#[cfg(CONFIG_KUNIT="y")]` before the module declaration.
+    let config_kunit = "#[cfg(CONFIG_KUNIT=\"y\")]".to_owned().parse().unwrap();
     tokens.insert(
         0,
         TokenTree::Group(Group::new(Delimiter::None, config_kunit)),
@@ -98,6 +98,8 @@ pub(crate) fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
     // ```
     let mut kunit_macros = "".to_owned();
     let mut test_cases = "".to_owned();
+    let mut assert_macros = "".to_owned();
+    let path = crate::helpers::file();
     for test in &tests {
         let kunit_wrapper_fn_name = format!("kunit_rust_wrapper_{test}");
         let kunit_wrapper = format!(
@@ -109,6 +111,27 @@ pub(crate) fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
             "    ::kernel::kunit::kunit_case(::kernel::c_str!(\"{test}\"), {kunit_wrapper_fn_name}),"
         )
         .unwrap();
+        writeln!(
+            assert_macros,
+            r#"
+/// Overrides the usual [`assert!`] macro with one that calls KUnit instead.
+#[allow(unused)]
+macro_rules! assert {{
+    ($cond:expr $(,)?) => {{{{
+        kernel::kunit_assert!("{test}", "{path}", 0, $cond);
+    }}}}
+}}
+
+/// Overrides the usual [`assert_eq!`] macro with one that calls KUnit instead.
+#[allow(unused)]
+macro_rules! assert_eq {{
+    ($left:expr, $right:expr $(,)?) => {{{{
+        kernel::kunit_assert_eq!("{test}", "{path}", 0, $left, $right);
+    }}}}
+}}
+        "#
+        )
+        .unwrap();
     }
 
     writeln!(kunit_macros).unwrap();
@@ -147,10 +170,12 @@ pub(crate) fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
         }
     }
 
-    let mut new_body = TokenStream::from_iter(new_body);
-    new_body.extend::<TokenStream>(kunit_macros.parse().unwrap());
+    let mut final_body = TokenStream::new();
+    final_body.extend::<TokenStream>(assert_macros.parse().unwrap());
+    final_body.extend(new_body);
+    final_body.extend::<TokenStream>(kunit_macros.parse().unwrap());
 
-    tokens.push(TokenTree::Group(Group::new(Delimiter::Brace, new_body)));
+    tokens.push(TokenTree::Group(Group::new(Delimiter::Brace, final_body)));
 
     tokens.into_iter().collect()
 }
diff --git a/rust/macros/lib.rs b/rust/macros/lib.rs
index b0b86fb9bfab..bb5df3b730fc 100644
--- a/rust/macros/lib.rs
+++ b/rust/macros/lib.rs
@@ -6,6 +6,11 @@
 // and thus add a dependency on `include/config/RUSTC_VERSION_TEXT`, which is
 // touched by Kconfig when the version string from the compiler changes.
 
+// Stable since Rust 1.88.0 under a different name, `proc_macro_span_file`,
+// which was added in Rust 1.88.0. This is why `cfg_attr` is used here, i.e.
+// to avoid depending on the full `proc_macro_span` on Rust >= 1.88.0.
+#![cfg_attr(not(CONFIG_RUSTC_HAS_SPAN_FILE), feature(proc_macro_span))]
+
 #[macro_use]
 mod quote;
 mod concat_idents;