forge/
progress.rs

1use alloy_primitives::map::HashMap;
2use chrono::Utc;
3use indicatif::{MultiProgress, ProgressBar};
4use parking_lot::Mutex;
5use std::{sync::Arc, time::Duration};
6
7/// State of [ProgressBar]s displayed for the given test run.
8/// Shows progress of all test suites matching filter.
9/// For each test within the test suite an individual progress bar is displayed.
10/// When a test suite completes, their progress is removed from overall progress and result summary
11/// is displayed.
12#[derive(Debug)]
13pub struct TestsProgressState {
14    /// Main [MultiProgress] instance showing progress for all test suites.
15    multi: MultiProgress,
16    /// Progress bar counting completed / remaining test suites.
17    overall_progress: ProgressBar,
18    /// Individual test suites progress.
19    suites_progress: HashMap<String, ProgressBar>,
20}
21
22impl TestsProgressState {
23    // Creates overall tests progress state.
24    pub fn new(suites_len: usize, threads_no: usize) -> Self {
25        let multi = MultiProgress::new();
26        let overall_progress = multi.add(ProgressBar::new(suites_len as u64));
27        overall_progress.set_style(
28            indicatif::ProgressStyle::with_template("{bar:40.cyan/blue} {pos:>7}/{len:7} {msg}")
29                .unwrap()
30                .progress_chars("##-"),
31        );
32        overall_progress.set_message(format!("completed (with {} threads)", threads_no as u64));
33        Self { multi, overall_progress, suites_progress: HashMap::default() }
34    }
35
36    /// Creates new test suite progress and add it to overall progress.
37    pub fn start_suite_progress(&mut self, suite_name: &String) {
38        let suite_progress = self.multi.add(ProgressBar::new_spinner());
39        suite_progress.set_style(
40            indicatif::ProgressStyle::with_template("{spinner} {wide_msg:.bold.dim}")
41                .unwrap()
42                .tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ "),
43        );
44        suite_progress.set_message(format!("{suite_name} "));
45        suite_progress.enable_steady_tick(Duration::from_millis(100));
46        self.suites_progress.insert(suite_name.to_owned(), suite_progress);
47    }
48
49    /// Prints suite result summary and removes it from overall progress.
50    pub fn end_suite_progress(&mut self, suite_name: &String, result_summary: String) {
51        if let Some(suite_progress) = self.suites_progress.remove(suite_name) {
52            self.multi.suspend(|| {
53                let _ = sh_println!("{suite_name}\n  ↪ {result_summary}");
54            });
55            suite_progress.finish_and_clear();
56            // Increment test progress bar to reflect completed test suite.
57            self.overall_progress.inc(1);
58        }
59    }
60
61    /// Creates progress entry for fuzz tests.
62    /// Set the prefix and total number of runs. Message is updated during execution with current
63    /// phase. Test progress is placed under test suite progress entry so all tests within suite
64    /// are grouped.
65    pub fn start_fuzz_progress(
66        &mut self,
67        suite_name: &str,
68        test_name: &String,
69        timeout: Option<u32>,
70        runs: u32,
71    ) -> Option<ProgressBar> {
72        if let Some(suite_progress) = self.suites_progress.get(suite_name) {
73            let fuzz_progress =
74                self.multi.insert_after(suite_progress, ProgressBar::new(runs as u64));
75            let template = if let Some(timeout) = timeout {
76                let ends_at = (Utc::now() + chrono::Duration::seconds(timeout.into()))
77                    .format("%H:%M:%S %Y-%m-%d")
78                    .to_string();
79                format!("    ↪ {{prefix:.bold.dim}}: [{{pos}}] Runs, ends at {ends_at} UTC {{msg}}")
80            } else {
81                "    ↪ {prefix:.bold.dim}: [{pos}/{len}] Runs {msg}".to_string()
82            };
83            fuzz_progress.set_style(
84                indicatif::ProgressStyle::with_template(&template).unwrap().tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ "),
85            );
86            fuzz_progress.set_prefix(test_name.to_string());
87            Some(fuzz_progress)
88        } else {
89            None
90        }
91    }
92
93    /// Removes overall test progress.
94    pub fn clear(&mut self) {
95        self.multi.clear().unwrap();
96    }
97}
98
99/// Cloneable wrapper around [TestsProgressState].
100#[derive(Debug, Clone)]
101pub struct TestsProgress {
102    pub inner: Arc<Mutex<TestsProgressState>>,
103}
104
105impl TestsProgress {
106    pub fn new(suites_len: usize, threads_no: usize) -> Self {
107        Self { inner: Arc::new(Mutex::new(TestsProgressState::new(suites_len, threads_no))) }
108    }
109}
110
111/// Helper function for creating fuzz test progress bar.
112pub fn start_fuzz_progress(
113    tests_progress: Option<&TestsProgress>,
114    suite_name: &str,
115    test_name: &String,
116    timeout: Option<u32>,
117    runs: u32,
118) -> Option<ProgressBar> {
119    if let Some(progress) = tests_progress {
120        progress.inner.lock().start_fuzz_progress(suite_name, test_name, timeout, runs)
121    } else {
122        None
123    }
124}