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