Skip to main content

writer/
test_buffer.rs

1// Copyright 2023 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use event_listener::Event;
6use futures::AsyncWrite;
7use std::io::Write;
8use std::pin::Pin;
9use std::sync::{Arc, Mutex};
10use std::task::{Context, Poll};
11
12/// Provides shared memory buffers (in the form of [`TestBuffer`]s for stdout
13/// and stderr that can be used with implementations of [`crate::ToolIO`] to
14/// test input and output behaviour at runtime.
15#[derive(Default, Debug)]
16pub struct TestBuffers {
17    pub stdout: TestBuffer,
18    pub stderr: TestBuffer,
19}
20
21/// Provides a shared memory buffer that can be cloned that can be used with
22/// implementations of [`crate::ToolIO`] to test input and output behaviour
23/// at runtime.
24#[derive(Clone, Default, Debug)]
25pub struct TestBuffer {
26    inner: Arc<Mutex<Vec<u8>>>,
27    event: Arc<Event>,
28}
29
30impl TestBuffers {
31    /// Destroy `self` and return the standard output and error buffers as byte
32    /// arrays.
33    pub fn into_stdio(self) -> (Vec<u8>, Vec<u8>) {
34        let stdout = self.stdout.into_inner();
35        let stderr = self.stderr.into_inner();
36        (stdout, stderr)
37    }
38
39    /// Destroy `self` and return the standard output and error buffers as strings.
40    pub fn into_strings(self) -> (String, String) {
41        let stdout = self.stdout.into_string();
42        let stderr = self.stderr.into_string();
43        (stdout, stderr)
44    }
45
46    /// Destroy `self` and return the standard output buffer as a byte
47    /// array.
48    pub fn into_stdout(self) -> Vec<u8> {
49        self.stdout.into_inner()
50    }
51
52    /// Destroy `self` and return the standard output buffer as a string.
53    pub fn into_stdout_str(self) -> String {
54        self.stdout.into_string()
55    }
56
57    /// Destroy `self` and return the standard error buffer as a byte
58    /// array.
59    pub fn into_stderr(self) -> Vec<u8> {
60        self.stderr.into_inner()
61    }
62
63    /// Destroy `self` and return the standard error buffer as a string.
64    pub fn into_stderr_str(self) -> String {
65        self.stderr.into_string()
66    }
67}
68
69impl TestBuffer {
70    /// Destroys self and returns the inner buffer as a vector.
71    pub fn into_inner(self) -> Vec<u8> {
72        std::mem::take(&mut *self.inner.lock().unwrap())
73    }
74
75    /// Destroys self and returns the inner buffer as a string.
76    pub fn into_string(self) -> String {
77        String::from_utf8(self.into_inner()).expect("Valid unicode on output string")
78    }
79
80    /// Waits for the buffer to be written to.  This will not return immediately if data is already
81    /// present so this should be called *after* calling `into_inner` or `into_string` and not
82    /// finding the desired results.
83    pub async fn wait_ready(&self) {
84        self.event.listen().await;
85    }
86}
87
88impl Write for TestBuffer {
89    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
90        self.event.notify(usize::MAX);
91        self.inner.lock().unwrap().write(buf)
92    }
93
94    fn flush(&mut self) -> std::io::Result<()> {
95        self.inner.lock().unwrap().flush()
96    }
97}
98
99impl AsyncWrite for TestBuffer {
100    fn poll_write(
101        self: Pin<&mut Self>,
102        _cx: &mut Context<'_>,
103        buf: &[u8],
104    ) -> Poll<std::io::Result<usize>> {
105        let size = self.inner.lock().unwrap().write(buf);
106        Poll::Ready(size)
107    }
108
109    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
110        match self.inner.lock().unwrap().flush() {
111            Ok(_) => Poll::Ready(Ok(())),
112            Err(e) => Poll::Ready(Err(e)),
113        }
114    }
115
116    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
117        Poll::Ready(Ok(()))
118    }
119}