Add a pipe stream for unit tests and a function for test timeouts.

src/network/mod.rs

@@ -42,3 +42,90 @@ pub struct Packet<Payload> {
     pub response: bool, // TODO: Not really required.
     pub payload: Payload,
 }
+
+#[cfg(test)]
+pub mod test_utils {
+    use futures::prelude::*;
+    use futures::sink::Sink;
+    use futures::stream::Stream;
+    use futures::task::Context;
+    use futures::task::Poll;
+    use std::pin::Pin;
+    use std::time::Duration;
+    use tokio::sync::mpsc;
+
+    pub struct Pipe {
+        sender: Pin<Box<mpsc::UnboundedSender<Result<Vec<u8>, PipeError>>>>,
+        receiver: Pin<Box<mpsc::UnboundedReceiver<Result<Vec<u8>, PipeError>>>>,
+    }
+
+    impl Pipe {
+        pub fn new() -> (Self, Self) {
+            let (send1, recv1) = mpsc::unbounded_channel();
+            let (send2, recv2) = mpsc::unbounded_channel();
+            (
+                Self {
+                    sender: Box::pin(send1),
+                    receiver: Box::pin(recv2),
+                },
+                Self {
+                    sender: Box::pin(send2),
+                    receiver: Box::pin(recv1),
+                },
+            )
+        }
+
+        pub fn inject_error(&mut self) {
+            self.sender.send(Err(PipeError::Injected)).unwrap();
+        }
+    }
+
+    impl Sink<Vec<u8>> for Pipe {
+        type Error = PipeError;
+
+        fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn start_send(self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
+            self.sender.send(Ok(item)).map_err(|_| PipeError::Closed)
+        }
+
+        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_close(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    impl Stream for Pipe {
+        type Item = Result<Vec<u8>, PipeError>;
+
+        fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
+            // Safe, as we will not move self_.
+            let self_ = unsafe { self.get_unchecked_mut() };
+
+            Pin::as_mut(&mut self_.receiver).poll_next(cx)
+        }
+    }
+
+    #[derive(Debug, PartialEq)]
+    pub enum PipeError {
+        Closed,
+        Injected,
+    }
+
+    pub async fn panic_after<T>(d: Duration, task: T) -> T::Output
+    where
+        T: Future + Send + 'static,
+        T::Output: Send + 'static,
+    {
+        let task = tokio::spawn(task);
+        match tokio::time::timeout(d, task).await {
+            Err(_) => panic!("function timed out ({:?})", d),
+            Ok(x) => x.unwrap(),
+        }
+    }
+}