improve tests

1 files changed, 56 insertions, 0 deletions

diff --git a/tests/helpers/mod.rs b/tests/helpers/mod.rs
new file mode 100644
index 0000000..4fee8df
--- /dev/null
+++ b/tests/helpers/mod.rs
@@ -0,0 +1,56 @@
+#![allow(dead_code)]
+
+use std::io::BufRead as _;
+
+pub struct Output<'a> {
+    pty: std::io::BufReader<&'a pty_process::blocking::Pty>,
+}
+
+impl<'a> Output<'a> {
+    fn new(pty: &'a pty_process::blocking::Pty) -> Self {
+        Self {
+            pty: std::io::BufReader::new(pty),
+        }
+    }
+}
+
+impl<'a> Iterator for Output<'a> {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let mut buf = vec![];
+        nix::unistd::alarm::set(5);
+        self.pty.read_until(b'\n', &mut buf).unwrap();
+        nix::unistd::alarm::cancel();
+        Some(std::string::String::from_utf8(buf).unwrap())
+    }
+}
+
+pub fn output(pty: &pty_process::blocking::Pty) -> Output<'_> {
+    Output::new(pty)
+}
+
+#[cfg(feature = "async")]
+pub fn output_async(
+    pty: &pty_process::Pty,
+) -> std::pin::Pin<Box<dyn futures::stream::Stream<Item = String> + '_>> {
+    use async_std::io::prelude::BufReadExt as _;
+    use futures::FutureExt as _;
+
+    let pty = async_std::io::BufReader::new(pty);
+    Box::pin(futures::stream::unfold(pty, |mut pty| async move {
+        Some((
+            async_std::future::timeout(
+                std::time::Duration::from_secs(5),
+                async {
+                    let mut buf = vec![];
+                    pty.read_until(b'\n', &mut buf).await.unwrap();
+                    std::string::String::from_utf8(buf).unwrap()
+                },
+            )
+            .map(|x| x.unwrap())
+            .await,
+            pty,
+        ))
+    }))
+}