From ec3daca5fd711ea7f3b2f6a67e8716114f49ef74 Mon Sep 17 00:00:00 2001 From: Zhang Junyang Date: Wed, 8 Nov 2023 23:59:23 +0800 Subject: [PATCH] Add ktest whitelist support --- .github/workflows/unit_test.yml | 2 +- Makefile | 19 +- README.md | 9 +- framework/jinux-frame/src/boot/kcmdline.rs | 17 +- framework/jinux-frame/src/boot/mod.rs | 29 +- framework/libs/ktest/src/lib.rs | 36 ++- framework/libs/ktest/src/path.rs | 284 +++++++++++++++++++ framework/libs/ktest/src/runner.rs | 115 ++++++-- framework/libs/ktest/src/tree.rs | 310 +++++++++++++++++++++ 9 files changed, 766 insertions(+), 55 deletions(-) create mode 100644 framework/libs/ktest/src/path.rs create mode 100644 framework/libs/ktest/src/tree.rs diff --git a/.github/workflows/unit_test.yml b/.github/workflows/unit_test.yml index ed6d65df9..e09a14ace 100644 --- a/.github/workflows/unit_test.yml +++ b/.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 diff --git a/Makefile b/Makefile index a63dc55ea..33a2a6bd2 100644 --- a/Makefile +++ b/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 \ diff --git a/README.md b/README.md index 770d666c9..7099bfe22 100644 --- a/README.md +++ b/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 diff --git a/framework/jinux-frame/src/boot/kcmdline.rs b/framework/jinux-frame/src/boot/kcmdline.rs index 0b943a7c9..c81d05b19 100644 --- a/framework/jinux-frame/src/boot/kcmdline.rs +++ b/framework/jinux-frame/src/boot/kcmdline.rs @@ -22,11 +22,17 @@ struct InitprocArgs { envp: Vec, } +#[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>, + module_args: BTreeMap>, } // 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> { + pub fn get_module_args(&self, module: &str) -> Option<&Vec> { 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 diff --git a/framework/jinux-frame/src/boot/mod.rs b/framework/jinux-frame/src/boot/mod.rs index 674c7dd4a..b0a0ab4e7 100644 --- a/framework/jinux-frame/src/boot/mod.rs +++ b/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::>(), + ); + } + } + _ => {} + } + } + } 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), } diff --git a/framework/libs/ktest/src/lib.rs b/framework/libs/ktest/src/lib.rs index 17f7011ea..5e2f1c24e 100644 --- a/framework/libs/ktest/src/lib.rs +++ b/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 -//! to pass a list of cfgs to the Makefile, e.g.: +//! variable to pass the cfgs to all rustc invocations. To run the tests, you simply +//! 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 -//! limited to crate names, enabling your imagination to configure running any subsets -//! of tests in any crates. And to ease development, `#[if_cfg_ktest]` is expanded to -//! a default conditional compilation setting: -//! `#[cfg(all(ktest, any(ktest = "all", ktest = #crate_name)))]` +//! Also, you can run a subset of tests by specifying the `KTEST_WHITELIST` variable. +//! This is achieved by a whitelist filter on the test name. +//! +//! ```bash +//! 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()) -} diff --git a/framework/libs/ktest/src/path.rs b/framework/libs/ktest/src/path.rs new file mode 100644 index 000000000..e62aa5eb0 --- /dev/null +++ b/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, +} + +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 { + 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 { + 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, + is_end: bool, +} + +impl SuffixTrie { + pub fn new() -> Self { + Self { + children: BTreeMap::new(), + is_end: false, + } + } + + pub fn from_paths>(paths: I) -> Self { + let mut t = Self::new(); + for i in paths { + t.insert(i.iter()); + } + t + } + + pub fn insert(&mut self, path: I) + where + I: DoubleEndedIterator, + P: Deref, + { + 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(&self, path: I) -> bool + where + I: DoubleEndedIterator, + P: Deref, + { + 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(&self, path: I) -> bool + where + I: DoubleEndedIterator, + P: Deref, + { + 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())); + } +} diff --git a/framework/libs/ktest/src/runner.rs b/framework/libs/ktest/src/runner.rs index 1490bf181..1f60ba962 100644 --- a/framework/libs/ktest/src/runner.rs +++ b/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(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( + print: &PrintFn, + catch_unwind: &CatchUnwindImpl, + whitelist: Option, +) -> KtestResult +where + PrintFn: Fn(core::fmt::Arguments), + PathsIter: Iterator, +{ + 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( + crate_: &KtestCrate, + print: &PrintFn, + catch_unwind: &CatchUnwindImpl, + whitelist: &Option, +) -> KtestResult where PrintFn: Fn(core::fmt::Arguments), { @@ -26,25 +79,41 @@ where } } - let (n, ktests) = crate::get_ktest_tests(); - print!("\nrunning {} tests\n\n", n); + let crate_name = crate_.name(); + 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 { - print!( - "[{}] test {}::{} ...", - test.info().package, - test.info().module_path, - test.info().fn_name - ); - match test.run(catch_unwind) { - Ok(()) => { - print!(" {}\n", "ok".green()); - passed += 1; + for module in crate_.iter() { + for test in module.iter() { + if let Some(trie) = whitelist { + let mut test_path = KtestPath::from(test.info().module_path); + test_path.push_back(test.info().fn_name); + if !trie.contains(test_path.iter()) { + filtered += 1; + continue; + } } - Err(e) => { - print!(" {}\n", "FAILED".red()); - failed_tests.push((test.clone(), e.clone())); + print!( + "test {}::{} ...", + 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 { diff --git a/framework/libs/ktest/src/tree.rs b/framework/libs/ktest/src/tree.rs new file mode 100644 index 000000000..913ebe0a5 --- /dev/null +++ b/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, + tests: Vec, +} + +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, +} + +impl FromIterator for KtestTree { + fn from_iter>(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.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 { + 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.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")); + } +}