chore: Racy buffering tests correction

.github/CODEOWNERS

@@ -75,4 +75,6 @@
 
 /src/event/merge.rs @MOZGIII
 
+/tests/buffering.rs @MOZGIII
+
 /website/ @binarylogic

src/test_util.rs

@@ -10,7 +10,7 @@ use std::fs::File;
 use std::io::Read;
 use std::iter;
 use std::mem;
-use std::net::SocketAddr;
+use std::net::{Shutdown, SocketAddr};
 use std::path::Path;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
@@ -66,7 +66,12 @@ pub fn send_lines(
                 .forward(out)
                 .and_then(|(_source, sink)| {
                     let socket = sink.into_inner().into_inner();
-                    tokio::io::shutdown(socket).map_err(|e| panic!("{:}", e))
+                    // In tokio 0.1 `AsyncWrite::shutdown` for `TcpStream` is a noop.
+                    // See https://docs.rs/tokio-tcp/0.1.4/src/tokio_tcp/stream.rs.html#917
+                    // Use `TcpStream::shutdown` instead - it actually does something.
+                    socket
+                        .shutdown(Shutdown::Both)
+                        .map_err(|e| panic!("{:}", e))
                 })
                 .map(|_| ())
         })
@@ -236,6 +241,15 @@ pub fn wait_for_tcp(addr: SocketAddr) {
     wait_for(|| std::net::TcpStream::connect(addr).is_ok())
 }
 
+pub fn wait_for_atomic_usize<T, F>(val: T, unblock: F)
+where
+    T: AsRef<AtomicUsize>,
+    F: Fn(usize) -> bool,
+{
+    let val = val.as_ref();
+    wait_for(|| unblock(val.load(Ordering::SeqCst)))
+}
+
 pub fn shutdown_on_idle(runtime: Runtime) {
     block_on(
         runtime
@@ -345,25 +359,25 @@ where
     }
 }
 
-pub struct Receiver {
-    handle: futures01::sync::oneshot::SpawnHandle<Vec<String>, ()>,
+pub struct Receiver<T> {
+    handle: futures01::sync::oneshot::SpawnHandle<Vec<T>, ()>,
     count: Arc<AtomicUsize>,
     trigger: Trigger,
     _runtime: Runtime,
 }
 
-impl Receiver {
+impl<T> Receiver<T> {
     pub fn count(&self) -> usize {
         self.count.load(Ordering::Relaxed)
     }
 
-    pub fn wait(self) -> Vec<String> {
+    pub fn wait(self) -> Vec<T> {
         self.trigger.cancel();
         self.handle.wait().unwrap()
     }
 }
 
-pub fn receive(addr: &SocketAddr) -> Receiver {
+pub fn receive(addr: &SocketAddr) -> Receiver<String> {
     let runtime = runtime();
 
     let listener = TcpListener::bind(addr).unwrap();
@@ -393,6 +407,35 @@ pub fn receive(addr: &SocketAddr) -> Receiver {
     }
 }
 
+pub fn receive_events<S>(stream: S) -> Receiver<Event>
+where
+    S: Stream<Item = Event> + Send + 'static,
+    <S as Stream>::Error: std::fmt::Debug,
+{
+    let runtime = runtime();
+
+    let count = Arc::new(AtomicUsize::new(0));
+    let count_clone = Arc::clone(&count);
+
+    let (trigger, tripwire) = Tripwire::new();
+
+    let events = stream
+        .take_until(tripwire)
+        .inspect(move |_| {
+            count_clone.fetch_add(1, Ordering::Relaxed);
+        })
+        .map_err(|e| panic!("{:?}", e))
+        .collect();
+
+    let handle = futures01::sync::oneshot::spawn(events, &runtime.executor());
+    Receiver {
+        handle,
+        count,
+        trigger,
+        _runtime: runtime,
+    }
+}
+
 pub struct CountReceiver {
     handle: futures01::sync::oneshot::SpawnHandle<usize, ()>,
     trigger: Trigger,