xacli_components/components/
progress.rs

1use std::{
2    io::{stdout, Write},
3    sync::mpsc::{self, Receiver, RecvTimeoutError, Sender},
4    time::Duration,
5};
6
7use crossterm::{
8    cursor, queue,
9    style::Print,
10    terminal::{self, ClearType},
11};
12
13use crate::Result;
14
15pub struct ProgressBar {
16    total: u64,
17    message: String,
18    rx: Receiver<ProgressEvent>,
19}
20
21pub struct ProgressHandle {
22    tx: Sender<ProgressEvent>,
23}
24
25pub enum ProgressEvent {
26    Increment(u64),
27    SetProgress(u64),
28    SetMessage(String),
29    Finish,
30}
31
32impl ProgressBar {
33    pub fn new(total: u64, message: impl Into<String>) -> (Self, ProgressHandle) {
34        let (tx, rx) = mpsc::channel();
35        let bar = Self {
36            total,
37            message: message.into(),
38            rx,
39        };
40        let handle = ProgressHandle { tx };
41        (bar, handle)
42    }
43
44    pub fn run(mut self) -> Result<()> {
45        let mut stdout = stdout();
46        let mut current = 0u64;
47
48        self.render(&mut stdout, current)?;
49
50        loop {
51            match self.rx.recv_timeout(Duration::from_millis(100)) {
52                Ok(ProgressEvent::Increment(n)) => {
53                    current = (current + n).min(self.total);
54                    self.render(&mut stdout, current)?;
55                }
56                Ok(ProgressEvent::SetProgress(n)) => {
57                    current = n.min(self.total);
58                    self.render(&mut stdout, current)?;
59                }
60                Ok(ProgressEvent::SetMessage(msg)) => {
61                    self.message = msg;
62                    self.render(&mut stdout, current)?;
63                }
64                Ok(ProgressEvent::Finish) => {
65                    current = self.total;
66                    self.render(&mut stdout, current)?;
67                    println!(); // 换行
68                    break;
69                }
70                Err(RecvTimeoutError::Timeout) => {
71                    // 定期刷新
72                    self.render(&mut stdout, current)?;
73                }
74                Err(RecvTimeoutError::Disconnected) => {
75                    println!(); // 换行
76                    break;
77                }
78            }
79        }
80
81        Ok(())
82    }
83
84    fn render(&self, stdout: &mut impl Write, current: u64) -> Result<()> {
85        let percent = if self.total > 0 {
86            (current as f64 / self.total as f64) * 100.0
87        } else {
88            0.0
89        };
90
91        let bar_width = 40;
92        let filled = if self.total > 0 {
93            ((current as f64 / self.total as f64) * bar_width as f64) as usize
94        } else {
95            0
96        };
97        let bar = format!("[{}{}]", "=".repeat(filled), " ".repeat(bar_width - filled));
98
99        queue!(
100            stdout,
101            cursor::MoveToColumn(0),
102            terminal::Clear(ClearType::CurrentLine),
103            Print(&self.message),
104            Print(" "),
105            Print(&bar),
106            Print(format!(" {}/{} ({:.1}%)", current, self.total, percent)),
107        )?;
108
109        stdout.flush()?;
110        Ok(())
111    }
112}
113
114impl ProgressHandle {
115    pub fn inc(&self, n: u64) {
116        let _ = self.tx.send(ProgressEvent::Increment(n));
117    }
118
119    pub fn set(&self, n: u64) {
120        let _ = self.tx.send(ProgressEvent::SetProgress(n));
121    }
122
123    pub fn set_message(&self, msg: impl Into<String>) {
124        let _ = self.tx.send(ProgressEvent::SetMessage(msg.into()));
125    }
126
127    pub fn finish(&self) {
128        let _ = self.tx.send(ProgressEvent::Finish);
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use std::{thread, time::Duration};
135
136    use super::*;
137
138    #[test]
139    fn test_progress_bar_new() {
140        let (bar, _handle) = ProgressBar::new(100, "Downloading");
141        assert_eq!(bar.total, 100);
142        assert_eq!(bar.message, "Downloading");
143    }
144
145    #[test]
146    fn test_progress_handle_inc() {
147        let (bar, handle) = ProgressBar::new(100, "Test");
148
149        handle.inc(10);
150        handle.inc(20);
151
152        // Check that events are received
153        let event1 = bar.rx.recv_timeout(Duration::from_millis(100)).unwrap();
154        assert!(matches!(event1, ProgressEvent::Increment(10)));
155
156        let event2 = bar.rx.recv_timeout(Duration::from_millis(100)).unwrap();
157        assert!(matches!(event2, ProgressEvent::Increment(20)));
158    }
159
160    #[test]
161    fn test_progress_handle_set() {
162        let (bar, handle) = ProgressBar::new(100, "Test");
163
164        handle.set(50);
165
166        let event = bar.rx.recv_timeout(Duration::from_millis(100)).unwrap();
167        assert!(matches!(event, ProgressEvent::SetProgress(50)));
168    }
169
170    #[test]
171    fn test_progress_handle_set_message() {
172        let (bar, handle) = ProgressBar::new(100, "Test");
173
174        handle.set_message("New message");
175
176        let event = bar.rx.recv_timeout(Duration::from_millis(100)).unwrap();
177        match event {
178            ProgressEvent::SetMessage(msg) => assert_eq!(msg, "New message"),
179            _ => panic!("Expected SetMessage event"),
180        }
181    }
182
183    #[test]
184    fn test_progress_handle_finish() {
185        let (bar, handle) = ProgressBar::new(100, "Test");
186
187        handle.finish();
188
189        let event = bar.rx.recv_timeout(Duration::from_millis(100)).unwrap();
190        assert!(matches!(event, ProgressEvent::Finish));
191    }
192
193    #[test]
194    fn test_render_zero_total() {
195        let (bar, _handle) = ProgressBar::new(0, "Test");
196        let mut output = Vec::new();
197
198        // render should handle zero total without division by zero
199        bar.render(&mut output, 0).unwrap();
200
201        // Output contains ANSI escape sequences, just verify it doesn't panic
202        assert!(!output.is_empty());
203    }
204
205    #[test]
206    fn test_render_with_progress() {
207        let (bar, _handle) = ProgressBar::new(100, "Loading");
208        let mut output = Vec::new();
209
210        bar.render(&mut output, 50).unwrap();
211
212        let output_str = String::from_utf8_lossy(&output);
213        assert!(output_str.contains("Loading"));
214        assert!(output_str.contains("50/100"));
215        assert!(output_str.contains("50.0%"));
216    }
217
218    #[test]
219    fn test_render_full_progress() {
220        let (bar, _handle) = ProgressBar::new(100, "Done");
221        let mut output = Vec::new();
222
223        bar.render(&mut output, 100).unwrap();
224
225        let output_str = String::from_utf8_lossy(&output);
226        assert!(output_str.contains("100/100"));
227        assert!(output_str.contains("100.0%"));
228    }
229
230    #[test]
231    fn test_progress_bar_run_with_finish() {
232        let (bar, handle) = ProgressBar::new(10, "Test");
233
234        // Spawn a thread to finish the progress bar
235        let finish_thread = thread::spawn(move || {
236            thread::sleep(Duration::from_millis(50));
237            handle.inc(5);
238            thread::sleep(Duration::from_millis(50));
239            handle.finish();
240        });
241
242        // Run should complete when finish is called
243        let result = bar.run();
244        assert!(result.is_ok());
245
246        finish_thread.join().unwrap();
247    }
248
249    #[test]
250    fn test_progress_bar_run_with_disconnect() {
251        let (bar, handle) = ProgressBar::new(10, "Test");
252
253        // Spawn a thread to drop the handle (disconnect)
254        let drop_thread = thread::spawn(move || {
255            thread::sleep(Duration::from_millis(50));
256            drop(handle);
257        });
258
259        // Run should complete when handle is dropped
260        let result = bar.run();
261        assert!(result.is_ok());
262
263        drop_thread.join().unwrap();
264    }
265
266    #[test]
267    fn test_progress_handle_after_disconnect() {
268        let (bar, handle) = ProgressBar::new(100, "Test");
269
270        // Drop the bar (receiver)
271        drop(bar);
272
273        // Sending should not panic, just silently fail
274        handle.inc(10);
275        handle.set(50);
276        handle.set_message("test");
277        handle.finish();
278    }
279
280    #[test]
281    fn test_progress_clamp_to_max() {
282        let (bar, handle) = ProgressBar::new(100, "Test");
283
284        // Set progress beyond total
285        handle.set(200);
286
287        let event = bar.rx.recv_timeout(Duration::from_millis(100)).unwrap();
288        // Event should still have the value 200, clamping happens in run()
289        assert!(matches!(event, ProgressEvent::SetProgress(200)));
290    }
291}