From 91c135a86ddefd0546fca128dc34ab5ad8efac1b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Zemanovi=C4=8D?= <tzemanovic@gmail.com>
Date: Mon, 31 Jan 2022 17:17:42 +0100
Subject: [PATCH] session: add `expect_eager` method

---
 src/session/async_session.rs | 51 ++++++++++++++++++++++++++++++++++++
 src/session/sync_session.rs  | 45 +++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/src/session/async_session.rs b/src/session/async_session.rs
index ef1ed1f..6223e95 100644
--- a/src/session/async_session.rs
+++ b/src/session/async_session.rs
@@ -80,6 +80,29 @@ impl<P, S: AsyncRead + Unpin> Session<P, S> {
         self.stream.expect(needle).await
     }
 
+    /// Expect waits until a pattern is matched. This continuously reads all the
+    /// available output, consumes it from the buffer and try to match on it
+    /// until a match is found.
+    ///
+    /// Typically, you'd want to use `expect_eager` over `expect` when the
+    /// command under test produces too much output to be processed lazily.
+    ///
+    /// If the method returns [Ok] it is guaranteed that at least 1 match was found.
+    ///
+    /// ```
+    /// # futures_lite::future::block_on(async {
+    /// let mut p = expectrl::spawn("echo 123").unwrap();
+    /// let m = p.expect_eager(expectrl::Regex("\\d+")).await.unwrap();
+    /// assert_eq!(m.first(), b"123");
+    /// # })
+    /// ```
+    ///
+    /// It return an error if timeout is reached.
+    /// You can specify a timeout value by [Session::set_expect_timeout] method.
+    pub async fn expect_eager<E: Needle>(&mut self, expect: E) -> Result<Found, Error> {
+        self.stream.expect_eager(needle).await
+    }
+
     /// Check checks if a pattern is matched.
     /// Returns empty found structure if nothing found.
     ///
@@ -327,6 +350,34 @@ impl<S: AsyncRead + Unpin> Stream<S> {
         }
     }
 
+    async fn expect_eager<N: Needle>(&mut self, needle: N) -> Result<Captures, Error> {
+        let start = time::Instant::now();
+        loop {
+            if let Some(timeout) = self.expect_timeout {
+                if start.elapsed() > timeout {
+                    return Err(Error::ExpectTimeout);
+                }
+            }
+
+            let eof = self.stream.read_available().await?;
+            let data = self.stream.get_available();
+
+            let found = expect.check(data, eof)?;
+            if !found.is_empty() {
+                let end_index = Found::right_most_index(&found);
+                let involved_bytes = data[..end_index].to_vec();
+                self.stream.consume_from_buffer(end_index);
+                return Ok(Found::new(involved_bytes, found));
+            } else if !data.is_empty() {
+                let data_len = data.len();
+                self.stream.consume_from_buffer(data_len);
+                continue;
+            } else if eof {
+                return Err(Error::Eof);
+            }
+        }
+    }
+
     /// Is matched checks if a pattern is matched.
     /// It doesn't consumes bytes from stream.
     async fn is_matched<E: Needle>(&mut self, needle: E) -> Result<bool, Error> {
diff --git a/src/session/sync_session.rs b/src/session/sync_session.rs
index 1232883..5deda20 100644
--- a/src/session/sync_session.rs
+++ b/src/session/sync_session.rs
@@ -132,6 +132,51 @@ impl<P, S: Read + NonBlocking> Session<P, S> {
         }
     }
 
+    /// Expect waits until a pattern is matched. This continuously reads all the
+    /// available output, consumes it from the buffer and try to match on it
+    /// until a match is found.
+    ///
+    /// Typically, you'd want to use `expect_eager` over `expect` when the
+    /// command under test produces too much output to be processed lazily.
+    ///
+    /// If the method returns [Ok] it is guaranteed that at least 1 match was found.
+    ///
+    /// ```
+    /// let mut p = expectrl::spawn("echo 123").unwrap();
+    /// let m = p.expect_eager(expectrl::Regex("\\d+")).unwrap();
+    /// assert_eq!(m.first(), b"123");
+    /// ```
+    ///
+    /// It returns an error if timeout is reached.
+    /// You can specify a timeout value by [Session::set_expect_timeout] method.
+    pub fn expect_eager<E: Needle>(&mut self, expect: E) -> Result<Captures, Error> {
+        let start = time::Instant::now();
+        loop {
+            if let Some(timeout) = self.expect_timeout {
+                if start.elapsed() > timeout {
+                    return Err(Error::ExpectTimeout);
+                }
+            }
+
+            let eof = self.stream.read_available()?;
+            let data = self.stream.get_available();
+
+            let found = expect.check(data, eof)?;
+            if !found.is_empty() {
+                let end_index = Captures::right_most_index(&found);
+                let involved_bytes = data[..end_index].to_vec();
+                self.stream.consume_available(end_index);
+                return Ok(Captures::new(involved_bytes, found));
+            } else if !data.is_empty() {
+                let data_len = data.len();
+                self.stream.consume_available(data_len);
+                continue;
+            } else if eof {
+                return Err(Error::Eof);
+            }
+        }
+    }
+
     /// Check verifies if a pattern is matched.
     /// Returns empty found structure if nothing found.
     ///