capture_stdio/
output_capture.rs

1//! Wrapper of `std::io::set_output_capture`
2
3use std::{
4    io::{set_output_capture, Error},
5    sync::{Arc, Mutex},
6};
7
8use crate::Capture;
9
10/// Wrap `std::io::set_output_capture`
11///
12/// You can use it to capture outputs of [print!] and [eprint!].
13pub struct OutputCapture {
14    local_stream: Arc<Mutex<Vec<u8>>>,
15    original: Option<Arc<Mutex<Vec<u8>>>>,
16    restored: bool,
17}
18
19impl Capture for OutputCapture {
20    fn capture() -> Result<Self, Error> {
21        let local_stream = Arc::new(Mutex::new(vec![]));
22        let original = set_output_capture(Some(local_stream.clone()));
23        Ok(Self {
24            local_stream,
25            original,
26            restored: false,
27        })
28    }
29
30    fn restore(&mut self) {
31        assert!(!self.restored, "You can't restore it twice");
32
33        set_output_capture(self.original.clone());
34        self.restored = true;
35    }
36}
37
38impl OutputCapture {
39    /// Get the captured output
40    pub fn get_output(&self) -> Arc<Mutex<Vec<u8>>> {
41        self.local_stream.clone()
42    }
43}
44
45impl Drop for OutputCapture {
46    fn drop(&mut self) {
47        if !self.restored {
48            self.restore();
49        }
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use crate::{output_capture::OutputCapture, Capture};
56
57    #[test]
58    fn test_output_capture() {
59        println!("This should not be captured");
60
61        let mut output_capture = OutputCapture::capture().unwrap();
62        println!("This should be captured");
63        output_capture.restore();
64
65        println!("This should not be captured");
66
67        let output =
68            String::from_utf8(output_capture.get_output().lock().unwrap().clone()).unwrap();
69
70        assert_eq!(output, "This should be captured\n");
71    }
72}
capture_stdio/output_capture.rs

capture_stdio/
output_capture.rs
