Merge pull request #10 from jayvdb/wait-for-write

Allow waiting for write

Author: lazy-bitfield

src/lib.rs

@@ -6,11 +6,23 @@ use std::cell::RefCell;
 use std::io::{Cursor, Error, ErrorKind, Read, Result, Write};
 use std::mem::swap;
 use std::rc::Rc;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
+use std::thread::sleep;
+use std::time;
 
 #[cfg(test)]
 mod tests;
 
+fn find_subsequence<T>(haystack: &[T], needle: &[T]) -> Option<usize>
+where
+    for<'a> &'a [T]: PartialEq,
+{
+    haystack
+        .windows(needle.len())
+        .position(|window| window == needle)
+}
+
 /// MockStream is Read+Write stream that stores the data written and provides the data to be read.
 #[derive(Clone)]
 pub struct MockStream {
@@ -37,6 +49,11 @@ impl MockStream {
         }
     }
 
+    /// Extract all bytes written by Write trait calls.
+    pub fn peek_bytes_written(&mut self) -> &Vec<u8> {
+        self.writer.get_ref()
+    }
+
     /// Extract all bytes written by Write trait calls.
     pub fn pop_bytes_written(&mut self) -> Vec<u8> {
         let mut result = Vec::new();
@@ -114,6 +131,8 @@ impl Write for SharedMockStream {
 #[derive(Clone, Default)]
 pub struct SyncMockStream {
     pimpl: Arc<Mutex<MockStream>>,
+    pub waiting_for_write: Arc<AtomicBool>,
+    pub expected_bytes: Vec<u8>,
 }
 
 impl SyncMockStream {
@@ -122,6 +141,12 @@ impl SyncMockStream {
         SyncMockStream::default()
     }
 
+    /// Block reads until expected bytes are written.
+    pub fn wait_for(&mut self, expected_bytes: &[u8]) {
+        self.expected_bytes = expected_bytes.to_vec();
+        self.waiting_for_write.store(true, Ordering::Relaxed);
+    }
+
     /// Extract all bytes written by Write trait calls.
     pub fn push_bytes_to_read(&mut self, bytes: &[u8]) {
         self.pimpl.lock().unwrap().push_bytes_to_read(bytes)
@@ -135,13 +160,27 @@ impl SyncMockStream {
 
 impl Read for SyncMockStream {
     fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
+        while self.waiting_for_write.load(Ordering::Relaxed) {
+            sleep(time::Duration::from_millis(10));
+        }
         self.pimpl.lock().unwrap().read(buf)
     }
 }
 
 impl Write for SyncMockStream {
     fn write(&mut self, buf: &[u8]) -> Result<usize> {
-        self.pimpl.lock().unwrap().write(buf)
+        let mut x = self.pimpl.lock().unwrap();
+        match x.write(buf) {
+            Ok(rv) => {
+                if self.waiting_for_write.load(Ordering::Relaxed)
+                    && find_subsequence(x.peek_bytes_written(), &self.expected_bytes).is_some()
+                {
+                    self.waiting_for_write.store(false, Ordering::Relaxed);
+                }
+                Ok(rv)
+            }
+            Err(rv) => Err(rv),
+        }
     }
 
     fn flush(&mut self) -> Result<()> {

src/tests.rs

@@ -177,3 +177,37 @@ fn test_sync_mock_stream() {
     assert_eq!(s.pop_bytes_written(), &[1, 2, 3, 4]);
     assert_eq!(read, &[5, 6, 7, 8]);
 }
+
+#[test]
+fn test_sync_mock_stream_wait_for() {
+    use std::thread;
+    use std::time::{Duration, SystemTime};
+    let mut s = SyncMockStream::new();
+
+    s.wait_for(&[3, 4]);
+
+    let mut s2 = s.clone();
+
+    // thread will write some bytes, and then read some bytes
+    s.push_bytes_to_read(&[5, 6, 7, 8]);
+    let read_thread = thread::spawn(move || {
+        let mut buf = Vec::new();
+        s2.read_to_end(&mut buf).unwrap();
+        let read_ts = SystemTime::now();
+        (buf, read_ts)
+    });
+
+    let write_thread = thread::spawn(move || {
+        thread::sleep(Duration::new(2, 0));
+        s.write_all(&[1, 2, 3, 4]).unwrap();
+        let write_ts = SystemTime::now();
+        (s.pop_bytes_written(), write_ts)
+    });
+
+    let (read, read_ts) = read_thread.join().unwrap();
+    let (write, write_ts) = write_thread.join().unwrap();
+
+    assert_eq!(write, &[1, 2, 3, 4]);
+    assert_eq!(read, &[5, 6, 7, 8]);
+    assert!(read_ts > write_ts);
+}