Add ktest whitelist support

.github/workflows/unit_test.yml

@@ -18,7 +18,7 @@ jobs:
 
       - name: Ktest Unit Test
         id: ktest_unit_test
-        run: make run KTEST=all ENABLE_KVM=0 RELEASE_MODE=1
+        run: make run KTEST=1 ENABLE_KVM=0 RELEASE_MODE=1
 
       - name: Usermode Unit test
         id: usermode_unit_test

Makefile

@@ -8,18 +8,22 @@ ENABLE_KVM ?= 1
 GDB_CLIENT ?= 0
 GDB_SERVER ?= 0
 INTEL_TDX ?= 0
-KTEST ?= none
+KTEST ?= 0
+KTEST_CRATES ?= all
+KTEST_WHITELIST ?=
 SKIP_GRUB_MENU ?= 1
 RELEASE_MODE ?= 0
 # End of setting up Make varaiables
 
-KERNEL_CMDLINE := SHELL="/bin/sh" LOGNAME="root" HOME="/" USER="root" PATH="/bin" init=/usr/bin/busybox -- sh -l
+KERNEL_CMDLINE := SHELL="/bin/sh" LOGNAME="root" HOME="/" USER="root" PATH="/bin" init=/usr/bin/busybox
+KERNEL_CMDLINE += ktest.whitelist="$(KTEST_WHITELIST)"
+INIT_CMDLINE := sh -l
 ifeq ($(AUTO_TEST), syscall)
 BUILD_SYSCALL_TEST := 1
-KERNEL_CMDLINE += /opt/syscall_test/run_syscall_test.sh
+INIT_CMDLINE += /opt/syscall_test/run_syscall_test.sh
 endif
 ifeq ($(AUTO_TEST), boot)
-KERNEL_CMDLINE += -c exit 0
+INIT_CMDLINE += -c exit 0
 endif
 
 CARGO_KBUILD_ARGS :=
@@ -31,7 +35,7 @@ CARGO_KBUILD_ARGS += --release
 CARGO_KRUN_ARGS += --release
 endif
 
-CARGO_KRUN_ARGS += -- '$(KERNEL_CMDLINE)'
+CARGO_KRUN_ARGS += -- '$(KERNEL_CMDLINE) -- $(INIT_CMDLINE)'
 CARGO_KRUN_ARGS += --boot-method="$(BOOT_METHOD)"
 CARGO_KRUN_ARGS += --boot-protocol="$(BOOT_PROTOCOL)"
 
@@ -61,9 +65,9 @@ CARGO_KBUILD_ARGS += --features intel_tdx
 CARGO_KRUN_ARGS += --features intel_tdx
 endif
 
-ifneq ($(KTEST), none)
+ifeq ($(KTEST), 1)
 comma := ,
-GLOBAL_RUSTC_FLAGS += --cfg ktest --cfg ktest=\"$(subst $(comma),\" --cfg ktest=\",$(KTEST))\"
+GLOBAL_RUSTC_FLAGS += --cfg ktest --cfg ktest=\"$(subst $(comma),\" --cfg ktest=\",$(KTEST_CRATES))\"
 endif
 
 ifeq ($(SKIP_GRUB_MENU), 1)
@@ -81,6 +85,7 @@ USERMODE_TESTABLE := \
     runner \
     framework/libs/align_ext \
     framework/libs/ktest \
+    framework/libs/ktest-proc-macro \
     services/libs/cpio-decoder \
     services/libs/int-to-c-enum \
     services/libs/int-to-c-enum/derive \

README.md

@@ -60,9 +60,8 @@ make run
 
 #### User mode unit tests
 
-Many of our crates does not require running on bare metal environment and can be tested through the standard Cargo testing framework. A specific list of which crates can be tested with `cargo test` is listed in the `[workspace.metadata.usermode_testable]` entry in the `Cargo.toml` file of the root workspace.
+Many of our crates does not require running on bare metal environment and can be tested through the standard Cargo testing framework. A specific list of which crates can be tested with `cargo test` is listed in the `Makefile` of the root workspace, and by using Make you can test all these crates together.
 
-There is a tool `./tools/test/run_tests.py` to run all the user mode tests, and can be invoked through Make.
 ```bash
 make test
 ```
@@ -73,12 +72,12 @@ Nevertheless, you could enter the directory of a specific crate and invoke `carg
 
 We can run unit tests in kernel mode for crates like `jinux-frame` or `jinux-std`. This is powered by our [ktest](framework/libs/ktest) framework.
 ```bash
-make run KTEST=all
+make run KTEST=1
 ```
 
-You could also specify tests in a crate or a subset of tests to run, as long as you defined them well using cfg.
+You could also specify tests in a crate or a subset of tests to run.
 ```bash
-make run KTEST=jinux-frame,jinux-std
+make run KTEST=1 KTEST_WHITELIST=failing_assertion,jinux_frame::test::expect_panic KTEST_CRATES=jinux-frame
 ```
 
 #### Component check

framework/jinux-frame/src/boot/kcmdline.rs

@@ -22,11 +22,17 @@ struct InitprocArgs {
     envp: Vec<CString>,
 }
 
+#[derive(PartialEq, Debug, Clone)]
+pub enum ModuleArg {
+    Arg(CString),
+    KeyVal(CString, CString),
+}
+
 /// The struct to store the parsed kernel command-line arguments.
 #[derive(Debug)]
 pub struct KCmdlineArg {
     initproc: InitprocArgs,
-    module_args: BTreeMap<String, Vec<CString>>,
+    module_args: BTreeMap<String, Vec<ModuleArg>>,
 }
 
 // Define get APIs.
@@ -44,7 +50,7 @@ impl KCmdlineArg {
         &self.initproc.envp
     }
     /// Get the argument vector of a kernel module.
-    pub fn get_module_args(&self, module: &str) -> Option<&Vec<CString>> {
+    pub fn get_module_args(&self, module: &str) -> Option<&Vec<ModuleArg>> {
         self.module_args.get(module)
     }
 }
@@ -121,9 +127,12 @@ impl From<&str> for KCmdlineArg {
             };
             if let Some(modname) = node {
                 let modarg = if let Some(v) = value {
-                    CString::new(option.to_string() + "=" + v).unwrap()
+                    ModuleArg::KeyVal(
+                        CString::new(option.to_string()).unwrap(),
+                        CString::new(v).unwrap(),
+                    )
                 } else {
-                    CString::new(option).unwrap()
+                    ModuleArg::Arg(CString::new(option).unwrap())
                 };
                 result
                     .module_args

framework/jinux-frame/src/boot/mod.rs

@@ -118,13 +118,38 @@ pub fn call_jinux_main() -> ! {
     #[cfg(ktest)]
     {
         use crate::arch::qemu::{exit_qemu, QemuExitCode};
-        use alloc::boxed::Box;
+        use alloc::{boxed::Box, string::ToString};
         use core::any::Any;
         crate::init();
         let fn_catch_unwind = &(unwinding::panic::catch_unwind::<(), fn()>
             as fn(fn()) -> Result<(), Box<(dyn Any + Send + 'static)>>);
+        // Parse the whitelist from the kernel command line.
+        let mut paths = None;
+        let args = kernel_cmdline().get_module_args("ktest");
+        if let Some(args) = args {
+            for options in args {
+                match options {
+                    kcmdline::ModuleArg::KeyVal(key, val) => {
+                        if key.to_str().unwrap() == "whitelist" && val.to_str().unwrap() != "" {
+                            let paths_str = val.to_str().unwrap();
+                            paths = Some(
+                                paths_str
+                                    .split(',')
+                                    .map(|s| s.to_string())
+                                    .collect::<Vec<_>>(),
+                            );
+                        }
+                    }
+                    _ => {}
+                }
+            }
+        }
         use ktest::runner::{run_ktests, KtestResult};
-        match run_ktests(crate::console::print, fn_catch_unwind) {
+        match run_ktests(
+            &crate::console::print,
+            fn_catch_unwind,
+            paths.map(|v| v.into_iter()),
+        ) {
             KtestResult::Ok => exit_qemu(QemuExitCode::Success),
             KtestResult::Failed => exit_qemu(QemuExitCode::Failed),
         }

framework/libs/ktest/src/lib.rs

@@ -56,18 +56,26 @@
 //!
 //! Rust cfg is used to control the compilation of the test module. In cooperation
 //! with the `ktest` framework, the Makefile will set the `RUSTFLAGS` environment
-//! variable to pass the cfgs to all rustc invocations. To run the tests, you need
+//! variable to pass the cfgs to all rustc invocations. To run the tests, you simply
-//! to pass a list of cfgs to the Makefile, e.g.:
+//! need to set a list of cfgs by specifying `KTEST=1` to the Makefile, e.g.:
 //!
 //! ```bash
-//! make run KTEST=jinux-frame,jinux-std,align_ext,tdx-guest
+//! make run KTEST=1
 //! ```
 //!
-//! It is flexible to specify the cfgs for running the tests. The cfg value is not
+//! Also, you can run a subset of tests by specifying the `KTEST_WHITELIST` variable.
-//! limited to crate names, enabling your imagination to configure running any subsets
+//! This is achieved by a whitelist filter on the test name.
-//! of tests in any crates. And to ease development, `#[if_cfg_ktest]` is expanded to
+//!
-//! a default conditional compilation setting:
+//! ```bash
-//! `#[cfg(all(ktest, any(ktest = "all", ktest = #crate_name)))]`
+//! make run KTEST=1 KTEST_WHITELIST=failing_assertion,jinux_frame::test::expect_panic
+//! ```
+//!
+//! `KTEST_CRATES` variable is used to specify in which crates the tests to be run.
+//! This is achieved by conditionally compiling the test module using the `#[cfg]`.
+//!
+//! ```bash
+//! make run KTEST=1 KTEST_CRATES=jinux-frame
+//! ``
 //!
 //! We support the `#[should_panic]` attribute just in the same way as the standard
 //! library do, but the implementation is quite slow currently. Use it with cautious.
@@ -76,10 +84,12 @@
 //! change.
 //!
 
-#![no_std]
+#![cfg_attr(not(test), no_std)]
 #![feature(panic_info_message)]
 
+pub mod path;
 pub mod runner;
+pub mod tree;
 
 extern crate alloc;
 use alloc::{boxed::Box, string::String};
@@ -110,7 +120,7 @@ pub enum KtestError {
     Unknown,
 }
 
-#[derive(Clone)]
+#[derive(Clone, PartialEq, Debug)]
 pub struct KtestItemInfo {
     pub module_path: &'static str,
     pub fn_name: &'static str,
@@ -120,7 +130,7 @@ pub struct KtestItemInfo {
     pub col: usize,
 }
 
-#[derive(Clone)]
+#[derive(Clone, PartialEq, Debug)]
 pub struct KtestItem {
     fn_: fn() -> (),
     should_panic: (bool, Option<&'static str>),
@@ -215,7 +225,3 @@ impl core::iter::Iterator for KtestIter {
         Some(ktest_item.clone())
     }
 }
-
-fn get_ktest_tests() -> (usize, KtestIter) {
-    (ktest_array!().len(), KtestIter::new())
-}

framework/libs/ktest/src/path.rs

@@ -0,0 +1,284 @@
+use alloc::{
+    collections::{vec_deque, BTreeMap, VecDeque},
+    string::{String, ToString},
+};
+use core::{fmt::Display, iter::zip, ops::Deref};
+
+pub type PathElement = String;
+
+pub type KtestPathIter<'a> = vec_deque::Iter<'a, PathElement>;
+
+#[derive(Debug)]
+pub struct KtestPath {
+    path: VecDeque<PathElement>,
+}
+
+impl From<&str> for KtestPath {
+    fn from(s: &str) -> Self {
+        let mut path = VecDeque::new();
+        for module in s.split("::") {
+            path.push_back(module.to_string());
+        }
+        Self { path }
+    }
+}
+
+impl KtestPath {
+    pub fn new() -> Self {
+        Self {
+            path: VecDeque::new(),
+        }
+    }
+
+    pub fn from(s: &str) -> Self {
+        Self {
+            path: s.split("::").map(PathElement::from).collect(),
+        }
+    }
+
+    pub fn push_back(&mut self, s: &str) {
+        self.path.push_back(PathElement::from(s));
+    }
+
+    pub fn pop_back(&mut self) -> Option<PathElement> {
+        self.path.pop_back()
+    }
+
+    pub fn push_front(&mut self, s: &str) {
+        self.path.push_front(PathElement::from(s))
+    }
+
+    pub fn pop_front(&mut self) -> Option<PathElement> {
+        self.path.pop_front()
+    }
+
+    pub fn len(&self) -> usize {
+        self.path.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.path.is_empty()
+    }
+
+    pub fn starts_with(&self, other: &Self) -> bool {
+        if self.path.len() < other.path.len() {
+            return false;
+        }
+        for (e1, e2) in zip(self.path.iter(), other.path.iter()) {
+            if e1 != e2 {
+                return false;
+            }
+        }
+        true
+    }
+
+    pub fn ends_with(&self, other: &Self) -> bool {
+        if self.path.len() < other.path.len() {
+            return false;
+        }
+        for (e1, e2) in zip(self.path.iter().rev(), other.path.iter().rev()) {
+            if e1 != e2 {
+                return false;
+            }
+        }
+        true
+    }
+
+    pub fn iter(&self) -> KtestPathIter {
+        self.path.iter()
+    }
+}
+
+impl Default for KtestPath {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Display for KtestPath {
+    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+        let mut first = true;
+        for e in self.path.iter() {
+            if first {
+                first = false;
+            } else {
+                write!(f, "::")?;
+            }
+            write!(f, "{}", e)?;
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod path_test {
+    use super::*;
+
+    #[test]
+    fn test_ktest_path() {
+        let mut path = KtestPath::new();
+        path.push_back("a");
+        path.push_back("b");
+        path.push_back("c");
+        assert_eq!(path.to_string(), "a::b::c");
+        assert_eq!(path.pop_back(), Some("c".to_string()));
+        assert_eq!(path.pop_back(), Some("b".to_string()));
+        assert_eq!(path.pop_back(), Some("a".to_string()));
+        assert_eq!(path.pop_back(), None);
+    }
+
+    #[test]
+    fn test_ktest_path_starts_with() {
+        let mut path = KtestPath::new();
+        path.push_back("a");
+        path.push_back("b");
+        path.push_back("c");
+        assert!(path.starts_with(&KtestPath::from("a")));
+        assert!(path.starts_with(&KtestPath::from("a::b")));
+        assert!(path.starts_with(&KtestPath::from("a::b::c")));
+        assert!(!path.starts_with(&KtestPath::from("a::b::c::d")));
+        assert!(!path.starts_with(&KtestPath::from("a::b::d")));
+        assert!(!path.starts_with(&KtestPath::from("a::d")));
+        assert!(!path.starts_with(&KtestPath::from("d")));
+    }
+
+    #[test]
+    fn test_ktest_path_ends_with() {
+        let mut path = KtestPath::new();
+        path.push_back("a");
+        path.push_back("b");
+        path.push_back("c");
+        assert!(path.ends_with(&KtestPath::from("c")));
+        assert!(path.ends_with(&KtestPath::from("b::c")));
+        assert!(path.ends_with(&KtestPath::from("a::b::c")));
+        assert!(!path.ends_with(&KtestPath::from("d::a::b::c")));
+        assert!(!path.ends_with(&KtestPath::from("a::b::d")));
+        assert!(!path.ends_with(&KtestPath::from("a::d")));
+        assert!(!path.ends_with(&KtestPath::from("d")));
+    }
+}
+
+#[derive(Debug)]
+pub struct SuffixTrie {
+    children: BTreeMap<PathElement, SuffixTrie>,
+    is_end: bool,
+}
+
+impl SuffixTrie {
+    pub fn new() -> Self {
+        Self {
+            children: BTreeMap::new(),
+            is_end: false,
+        }
+    }
+
+    pub fn from_paths<I: IntoIterator<Item = KtestPath>>(paths: I) -> Self {
+        let mut t = Self::new();
+        for i in paths {
+            t.insert(i.iter());
+        }
+        t
+    }
+
+    pub fn insert<I, P>(&mut self, path: I)
+    where
+        I: DoubleEndedIterator<Item = P>,
+        P: Deref<Target = PathElement>,
+    {
+        let mut cur = self;
+        for e in path.into_iter().rev() {
+            cur = cur.children.entry(e.clone()).or_default();
+        }
+        cur.is_end = true;
+    }
+
+    /// Find if there is a perfect match in this suffix trie.
+    pub fn matches<I, P>(&self, path: I) -> bool
+    where
+        I: DoubleEndedIterator<Item = P>,
+        P: Deref<Target = PathElement>,
+    {
+        let mut cur = self;
+        for e in path.into_iter().rev() {
+            if let Some(next) = cur.children.get(&*e) {
+                cur = next;
+            } else {
+                return false;
+            }
+        }
+        cur.is_end
+    }
+
+    /// Find if any suffix of the path exists in the suffix trie.
+    pub fn contains<I, P>(&self, path: I) -> bool
+    where
+        I: DoubleEndedIterator<Item = P>,
+        P: Deref<Target = PathElement>,
+    {
+        let mut cur = self;
+        for e in path.into_iter().rev() {
+            if let Some(next) = cur.children.get(&*e) {
+                cur = next;
+                if cur.is_end {
+                    return true;
+                }
+            } else {
+                return false;
+            }
+        }
+        false
+    }
+}
+
+impl Default for SuffixTrie {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod suffix_trie_test {
+    use super::*;
+
+    static TEST_PATHS: &[&str] = &[
+        "a::b::c::d",
+        "a::b::c::e",
+        "a::b::d::e",
+        "a::b::f::g",
+        "h::i::j::k",
+        "l::m::n",
+        "m::n",
+    ];
+
+    #[test]
+    fn test_contains() {
+        let trie = SuffixTrie::from_paths(TEST_PATHS.iter().map(|&s| KtestPath::from(s)));
+
+        assert!(trie.contains(KtestPath::from("e::f::g::a::b::c::d").iter()));
+        assert!(trie.contains(KtestPath::from("e::f::g::a::b::f::g").iter()));
+        assert!(trie.contains(KtestPath::from("h::i::j::l::m::n").iter()));
+        assert!(trie.contains(KtestPath::from("l::m::n").iter()));
+
+        assert!(!trie.contains(KtestPath::from("a::b::c").iter()));
+        assert!(!trie.contains(KtestPath::from("b::c::d").iter()));
+        assert!(!trie.contains(KtestPath::from("a::b::f").iter()));
+        assert!(!trie.contains(KtestPath::from("i::j").iter()));
+        assert!(!trie.contains(KtestPath::from("h::i::j::l::n").iter()));
+        assert!(!trie.contains(KtestPath::from("n").iter()));
+    }
+
+    #[test]
+    fn test_matches() {
+        let trie = SuffixTrie::from_paths(TEST_PATHS.iter().map(|&s| KtestPath::from(s)));
+
+        assert!(trie.matches(KtestPath::from("a::b::c::d").iter()));
+        assert!(trie.matches(KtestPath::from("a::b::c::e").iter()));
+        assert!(trie.matches(KtestPath::from("l::m::n").iter()));
+        assert!(trie.matches(KtestPath::from("m::n").iter()));
+
+        assert!(!trie.matches(KtestPath::from("a::b::d").iter()));
+        assert!(!trie.matches(KtestPath::from("b::c::e").iter()));
+        assert!(!trie.matches(KtestPath::from("l::n::h::k::j::i").iter()));
+        assert!(!trie.matches(KtestPath::from("n").iter()));
+    }
+}

framework/libs/ktest/src/runner.rs

@@ -1,6 +1,13 @@
-use crate::{CatchUnwindImpl, KtestError, KtestItem};
+//! Test runner enabling control over the tests.
+//!
 
-use alloc::vec::Vec;
+use crate::{
+    path::{KtestPath, SuffixTrie},
+    tree::{KtestCrate, KtestTree},
+    CatchUnwindImpl, KtestError, KtestItem, KtestIter,
+};
+
+use alloc::{string::String, vec::Vec};
 use core::format_args;
 
 use owo_colors::OwoColorize;
@@ -12,11 +19,57 @@ pub enum KtestResult {
 
 /// Run all the tests registered by `#[ktest]` in the `.ktest_array` section.
 ///
-/// Need to provide a print function to print the test result, and a `catch_unwind`
+/// Need to provide a print function `print` to print the test result, and a `catch_unwind`
 /// implementation to catch the panic.
 ///
+/// The `whitelist` argument is optional. If it is `None`, all tests compiled will be run.
+/// If it is `Some`, only the tests whose test path being the suffix of any paths in the whitelist
+/// will be run.
+///
 /// Returns the test result interpreted as `ok` or `FAILED`.
-pub fn run_ktests<PrintFn>(print: PrintFn, catch_unwind: &CatchUnwindImpl) -> KtestResult
+///
+/// If a test inside a crate fails, the test runner will continue to run the rest of the tests
+/// inside the crate. But the tests in the following crates will not be run.
+pub fn run_ktests<PrintFn, PathsIter>(
+    print: &PrintFn,
+    catch_unwind: &CatchUnwindImpl,
+    whitelist: Option<PathsIter>,
+) -> KtestResult
+where
+    PrintFn: Fn(core::fmt::Arguments),
+    PathsIter: Iterator<Item = String>,
+{
+    macro_rules! print {
+        ($fmt: literal $(, $($arg: tt)+)?) => {
+            print(format_args!($fmt $(, $($arg)+)?))
+        }
+    }
+
+    let whitelist_trie =
+        whitelist.map(|paths| SuffixTrie::from_paths(paths.map(|p| KtestPath::from(&p))));
+
+    let tree = KtestTree::from_iter(KtestIter::new());
+    print!(
+        "\n[ktest runner] running {} tests in {} crates\n",
+        tree.nr_tot_tests(),
+        tree.nr_tot_crates()
+    );
+    for crate_ in tree.iter() {
+        match run_crate_ktests(crate_, print, catch_unwind, &whitelist_trie) {
+            KtestResult::Ok => {}
+            KtestResult::Failed => return KtestResult::Failed,
+        }
+    }
+    print!("\n[ktest runner] All crates tested.\n");
+    KtestResult::Ok
+}
+
+fn run_crate_ktests<PrintFn>(
+    crate_: &KtestCrate,
+    print: &PrintFn,
+    catch_unwind: &CatchUnwindImpl,
+    whitelist: &Option<SuffixTrie>,
+) -> KtestResult
 where
     PrintFn: Fn(core::fmt::Arguments),
 {
@@ -26,25 +79,41 @@ where
         }
     }
 
-    let (n, ktests) = crate::get_ktest_tests();
+    let crate_name = crate_.name();
-    print!("\nrunning {} tests\n\n", n);
+    print!(
+        "\nrunning {} tests in crate \"{}\"\n\n",
+        crate_.nr_tot_tests(),
+        crate_name
+    );
+
     let mut passed: usize = 0;
+    let mut filtered: usize = 0;
     let mut failed_tests: Vec<(KtestItem, KtestError)> = Vec::new();
-    for test in ktests {
+    for module in crate_.iter() {
-        print!(
+        for test in module.iter() {
-            "[{}] test {}::{} ...",
+            if let Some(trie) = whitelist {
-            test.info().package,
+                let mut test_path = KtestPath::from(test.info().module_path);
-            test.info().module_path,
+                test_path.push_back(test.info().fn_name);
-            test.info().fn_name
+                if !trie.contains(test_path.iter()) {
-        );
+                    filtered += 1;
-        match test.run(catch_unwind) {
+                    continue;
-            Ok(()) => {
+                }
-                print!(" {}\n", "ok".green());
-                passed += 1;
             }
-            Err(e) => {
+            print!(
-                print!(" {}\n", "FAILED".red());
+                "test {}::{} ...",
-                failed_tests.push((test.clone(), e.clone()));
+                test.info().module_path,
+                test.info().fn_name
+            );
+            debug_assert_eq!(test.info().package, crate_name);
+            match test.run(catch_unwind) {
+                Ok(()) => {
+                    print!(" {}\n", "ok".green());
+                    passed += 1;
+                }
+                Err(e) => {
+                    print!(" {}\n", "FAILED".red());
+                    failed_tests.push((test.clone(), e.clone()));
+                }
             }
         }
     }
@@ -54,7 +123,11 @@ where
     } else {
         print!("\ntest result: {}.", "FAILED".red());
     }
-    print!(" {} passed; {} failed.\n", passed, failed);
+    print!(
+        " {} passed; {} failed; {} filtered out.\n",
+        passed, failed, filtered
+    );
+    assert!(passed + failed + filtered == crate_.nr_tot_tests());
     if failed > 0 {
         print!("\nfailures:\n\n");
         for (t, e) in failed_tests {

framework/libs/ktest/src/tree.rs

@@ -0,0 +1,310 @@
+//! The source module tree of ktests.
+//!
+//! In the `KtestTree`, the root is abstract, and the children of the root are the
+//! crates. The leaves are the test functions. Nodes other than the root and the
+//! leaves are modules.
+//!
+
+use alloc::{
+    collections::{btree_map, BTreeMap},
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
+use core::iter::{FromIterator, Iterator};
+
+use crate::{
+    path::{KtestPath, PathElement},
+    KtestItem,
+};
+
+#[derive(Debug)]
+pub struct KtestModule {
+    nr_tot_tests: usize,
+    name: PathElement,
+    children: BTreeMap<PathElement, KtestModule>,
+    tests: Vec<KtestItem>,
+}
+
+impl KtestModule {
+    pub fn nr_this_tests(&self) -> usize {
+        self.tests.len()
+    }
+
+    pub fn nr_tot_tests(&self) -> usize {
+        self.nr_tot_tests
+    }
+
+    pub fn name(&self) -> &PathElement {
+        &self.name
+    }
+
+    fn insert(&mut self, module_path: &mut KtestPath, test: KtestItem) {
+        self.nr_tot_tests += 1;
+        if module_path.is_empty() {
+            self.tests.push(test);
+        } else {
+            let module_name = module_path.pop_front().unwrap();
+            let node = self.children.entry(module_name.clone()).or_insert(Self {
+                nr_tot_tests: 0,
+                name: module_name,
+                children: BTreeMap::new(),
+                tests: Vec::new(),
+            });
+            node.nr_tot_tests += 1;
+            node.insert(module_path, test);
+        }
+    }
+
+    fn new(name: PathElement) -> Self {
+        Self {
+            nr_tot_tests: 0,
+            name,
+            children: BTreeMap::new(),
+            tests: Vec::new(),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct KtestCrate {
+    // Crate behaves just like modules, which can own it's children and tests.
+    // But the iterator it provides will only iterate over the modules not tests.
+    root_module: KtestModule,
+}
+
+impl KtestCrate {
+    pub fn nr_tot_tests(&self) -> usize {
+        self.root_module.nr_tot_tests()
+    }
+
+    pub fn name(&self) -> &str {
+        self.root_module.name()
+    }
+}
+
+pub struct KtestTree {
+    nr_tot_tests: usize,
+    crates: BTreeMap<String, KtestCrate>,
+}
+
+impl FromIterator<KtestItem> for KtestTree {
+    fn from_iter<I: IntoIterator<Item = KtestItem>>(iter: I) -> Self {
+        let mut tree = Self::new();
+        for test in iter {
+            tree.add_ktest(test);
+        }
+        tree
+    }
+}
+
+impl KtestTree {
+    pub fn new() -> Self {
+        Self {
+            nr_tot_tests: 0,
+            crates: BTreeMap::new(),
+        }
+    }
+
+    pub fn add_ktest(&mut self, test: KtestItem) {
+        self.nr_tot_tests += 1;
+        let package = test.info().package.to_string();
+        let module_path = test.info().module_path;
+        let node = self.crates.entry(package.clone()).or_insert(KtestCrate {
+            root_module: KtestModule::new(PathElement::from(package)),
+        });
+        node.root_module
+            .insert(&mut KtestPath::from(module_path), test);
+    }
+
+    pub fn nr_tot_tests(&self) -> usize {
+        self.nr_tot_tests
+    }
+
+    pub fn nr_tot_crates(&self) -> usize {
+        self.crates.len()
+    }
+}
+
+impl Default for KtestTree {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// The `KtestTreeIter` will iterate over all crates. Yeilding `KtestCrate`s.
+pub struct KtestTreeIter<'a> {
+    crate_iter: btree_map::Iter<'a, String, KtestCrate>,
+}
+
+impl KtestTree {
+    pub fn iter(&self) -> KtestTreeIter<'_> {
+        KtestTreeIter {
+            crate_iter: self.crates.iter(),
+        }
+    }
+}
+
+impl<'a> Iterator for KtestTreeIter<'a> {
+    type Item = &'a KtestCrate;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.crate_iter.next().map(|(_, v)| v)
+    }
+}
+
+type CrateChildrenIter<'a> = btree_map::Iter<'a, String, KtestModule>;
+
+/// The `KtestCrateIter` will iterate over all modules in a crate. Yeilding `KtestModule`s.
+/// The iterator will return modules in the depth-first-search order of the module tree.
+pub struct KtestCrateIter<'a> {
+    path: Vec<(&'a KtestModule, CrateChildrenIter<'a>)>,
+}
+
+impl KtestCrate {
+    pub fn iter(&self) -> KtestCrateIter<'_> {
+        KtestCrateIter {
+            path: vec![(&self.root_module, self.root_module.children.iter())],
+        }
+    }
+}
+
+impl<'a> Iterator for KtestCrateIter<'a> {
+    type Item = &'a KtestModule;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let next_module = loop {
+            let Some(last) = self.path.last_mut() else {
+                break None;
+            };
+            if let Some((_, next_module)) = last.1.next() {
+                break Some(next_module);
+            }
+            let (_, _) = self.path.pop().unwrap();
+        };
+        if let Some(next_module) = next_module {
+            self.path.push((next_module, next_module.children.iter()));
+            Some(next_module)
+        } else {
+            None
+        }
+    }
+}
+
+/// The `KtestModuleIter` will iterate over all tests in a crate. Yeilding `KtestItem`s.
+pub struct KtestModuleIter<'a> {
+    test_iter: core::slice::Iter<'a, KtestItem>,
+}
+
+impl KtestModule {
+    pub fn iter(&self) -> KtestModuleIter<'_> {
+        KtestModuleIter {
+            test_iter: self.tests.iter(),
+        }
+    }
+}
+
+impl<'a> Iterator for KtestModuleIter<'a> {
+    type Item = &'a KtestItem;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.test_iter.next()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    macro_rules! gen_test_case {
+        () => {{
+            fn dummy_fn() {
+                ()
+            }
+            let mut tree = KtestTree::new();
+            let new = |m: &'static str, f: &'static str, p: &'static str| {
+                KtestItem::new(
+                    dummy_fn,
+                    (false, None),
+                    crate::KtestItemInfo {
+                        module_path: m,
+                        fn_name: f,
+                        package: p,
+                        source: "unrelated",
+                        line: 0,
+                        col: 0,
+                    },
+                )
+            };
+            tree.add_ktest(new("crate1::mod1::mod2", "test21", "crate1"));
+            tree.add_ktest(new("crate1::mod1", "test11", "crate1"));
+            tree.add_ktest(new("crate1::mod1::mod2", "test22", "crate1"));
+            tree.add_ktest(new("crate1::mod1::mod2", "test23", "crate1"));
+            tree.add_ktest(new("crate1::mod1::mod3", "test31", "crate1"));
+            tree.add_ktest(new("crate1::mod1::mod3::mod4", "test41", "crate1"));
+            tree.add_ktest(new("crate2::mod1::mod2", "test2", "crate2"));
+            tree.add_ktest(new("crate2::mod1", "test1", "crate2"));
+            tree.add_ktest(new("crate2::mod1::mod2", "test3", "crate2"));
+            tree
+        }};
+    }
+
+    #[test]
+    fn test_tree_iter() {
+        let tree = gen_test_case!();
+        let mut iter = tree.iter();
+        let c1 = iter.next().unwrap();
+        assert_eq!(c1.name(), "crate1");
+        let c2 = iter.next().unwrap();
+        assert_eq!(c2.name(), "crate2");
+        assert!(iter.next().is_none());
+    }
+
+    #[test]
+    fn test_crate_iter() {
+        let tree = gen_test_case!();
+        for crate_ in tree.iter() {
+            if crate_.name() == "crate1" {
+                let mut len = 0;
+                for module in crate_.iter() {
+                    len += 1;
+                    let modules = ["crate1", "mod1", "mod2", "mod3", "mod4"];
+                    assert!(modules.contains(&module.name().as_str()));
+                }
+                assert_eq!(len, 5);
+            } else if crate_.name() == "crate2" {
+                let mut len = 0;
+                for module in crate_.iter() {
+                    len += 1;
+                    let modules = ["crate2", "mod1", "mod2"];
+                    assert!(modules.contains(&module.name().as_str()));
+                }
+                assert_eq!(len, 3);
+            }
+        }
+    }
+
+    #[test]
+    fn test_module_iter() {
+        let tree = gen_test_case!();
+        let mut collection = Vec::<&KtestItem>::new();
+        for crate_ in tree.iter() {
+            for mov in crate_.iter() {
+                let module = mov;
+                for test in module.iter() {
+                    collection.push(&test);
+                }
+            }
+        }
+        assert_eq!(collection.len(), 9);
+        assert!(collection.iter().any(|t| t.info().fn_name == "test1"));
+        assert!(collection.iter().any(|t| t.info().fn_name == "test2"));
+        assert!(collection.iter().any(|t| t.info().fn_name == "test3"));
+        assert!(collection.iter().any(|t| t.info().fn_name == "test11"));
+        assert!(collection.iter().any(|t| t.info().fn_name == "test21"));
+        assert!(collection.iter().any(|t| t.info().fn_name == "test22"));
+        assert!(collection.iter().any(|t| t.info().fn_name == "test23"));
+        assert!(collection.iter().any(|t| t.info().fn_name == "test31"));
+        assert!(collection.iter().any(|t| t.info().fn_name == "test41"));
+    }
+}