Add test infrastructure to work directly with frames (#56)

This adds a `Codec` based testing API. This is a bit less annoying than writing
at the raw H2 wire protocol level...

tests/support/src/assert.rs

@@ -0,0 +1,50 @@
+
+#[macro_export]
+macro_rules! assert_closed {
+    ($transport:expr) => {{
+        assert_eq!($transport.poll().unwrap(), None.into());
+    }}
+}
+
+#[macro_export]
+macro_rules! assert_ping {
+    ($frame:expr) => {{
+        match $frame {
+            ::h2::frame::Frame::Ping(v) => v,
+            f => panic!("expected PING; actual={:?}", f),
+        }
+    }}
+}
+
+#[macro_export]
+macro_rules! assert_settings {
+    ($frame:expr) => {{
+        match $frame {
+            ::h2::frame::Frame::Settings(v) => v,
+            f => panic!("expected SETTINGS; actual={:?}", f),
+        }
+    }}
+}
+
+#[macro_export]
+macro_rules! poll_err {
+    ($transport:expr) => {{
+        match $transport.poll() {
+            Err(e) => e,
+            frame => panic!("expected error; actual={:?}", frame),
+        }
+    }}
+}
+
+#[macro_export]
+macro_rules! poll_data {
+    ($transport:expr) => {{
+        use h2::frame::Frame;
+        use futures::Async;
+
+        match $transport.poll() {
+            Ok(Async::Ready(Some(Frame::Data(frame)))) => frame,
+            frame => panic!("expected data frame; actual={:?}", frame),
+        }
+    }}
+}

tests/support/src/codec.rs

@@ -1,74 +0,0 @@
-#[macro_export]
-macro_rules! assert_closed {
-    ($transport:expr) => {{
-        assert_eq!($transport.poll().unwrap(), None.into());
-    }}
-}
-
-#[macro_export]
-macro_rules! poll_err {
-    ($transport:expr) => {{
-        match $transport.poll() {
-            Err(e) => e,
-            frame => panic!("expected error; actual={:?}", frame),
-        }
-    }}
-}
-
-#[macro_export]
-macro_rules! poll_data {
-    ($transport:expr) => {{
-        use h2::frame::Frame;
-        use futures::Async;
-
-        match $transport.poll() {
-            Ok(Async::Ready(Some(Frame::Data(frame)))) => frame,
-            frame => panic!("expected data frame; actual={:?}", frame),
-        }
-    }}
-}
-
-#[macro_export]
-macro_rules! raw_codec {
-    (
-        $(
-            $fn:ident => [$($chunk:expr,)+];
-        )*
-    ) => {{
-        let mut b = $crate::mock_io::Builder::new();
-
-        $({
-            let mut chunk = vec![];
-
-            $(
-                $crate::codec::Chunk::push(&$chunk, &mut chunk);
-            )+
-
-            b.$fn(&chunk[..]);
-        })*
-
-        $crate::Codec::new(b.build())
-    }}
-}
-
-pub trait Chunk {
-    fn push(&self, dst: &mut Vec<u8>);
-}
-
-impl Chunk for u8 {
-    fn push(&self, dst: &mut Vec<u8>) {
-        dst.push(*self);
-    }
-}
-
-impl<'a> Chunk for &'a [u8] {
-    fn push(&self, dst: &mut Vec<u8>) {
-        dst.extend(*self)
-    }
-}
-
-impl<'a> Chunk for &'a str {
-    fn push(&self, dst: &mut Vec<u8>) {
-        dst.extend(self.as_bytes())
-    }
-}

tests/support/src/future_ext.rs

@@ -0,0 +1,32 @@
+use futures::{Future, Poll};
+
+use std::fmt;
+
+pub trait FutureExt: Future {
+    fn unwrap(self) -> Unwrap<Self>
+        where Self: Sized,
+              Self::Error: fmt::Debug,
+    {
+        Unwrap { inner: self }
+    }
+}
+
+pub struct Unwrap<T> {
+    inner: T,
+}
+
+impl<T: Future> FutureExt for T {
+}
+
+impl<T> Future for Unwrap<T>
+    where T: Future,
+          T::Item: fmt::Debug,
+          T::Error: fmt::Debug,
+{
+    type Item = T::Item;
+    type Error = ();
+
+    fn poll(&mut self) -> Poll<T::Item, ()> {
+        Ok(self.inner.poll().unwrap())
+    }
+}

tests/support/src/lib.rs

@@ -3,93 +3,28 @@
 pub extern crate bytes;
 pub extern crate h2;
 pub extern crate http;
+
+#[macro_use]
 pub extern crate tokio_io;
 pub extern crate futures;
 pub extern crate mock_io;
 pub extern crate env_logger;
 
 #[macro_use]
-pub mod codec;
+mod assert;
 
-pub use self::futures::{
-    Future,
-    Sink,
-    Stream,
-};
-pub use self::futures::future::poll_fn;
+#[macro_use]
+pub mod raw;
 
-pub use self::http::{
-    request,
-    response,
-    Request,
-    Response,
-    Method,
-    HeaderMap,
-    StatusCode,
-};
+pub mod prelude;
+pub mod mock;
 
-pub use self::h2::*;
-pub use self::h2::client::{self, Client};
-pub use self::h2::server::{self, Server};
+mod future_ext;
 
+pub use future_ext::{FutureExt, Unwrap};
+
+// This is our test Codec type
 pub type Codec<T> = h2::Codec<T, ::std::io::Cursor<::bytes::Bytes>>;
 
-pub use self::bytes::{
-    Buf,
-    BufMut,
-    Bytes,
-    BytesMut,
-    IntoBuf,
-};
-
-pub use std::time::Duration;
-
-use tokio_io::{AsyncRead, AsyncWrite};
-
-pub trait MockH2 {
-    fn handshake(&mut self) -> &mut Self;
-}
-
-impl MockH2 for mock_io::Builder {
-    fn handshake(&mut self) -> &mut Self {
-        self.write(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")
-            // Settings frame
-            .write(frames::SETTINGS)
-            .read(frames::SETTINGS)
-            .read(frames::SETTINGS_ACK)
-    }
-}
-
-pub trait ClientExt {
-    fn run<F: Future>(&mut self, f: F) -> Result<F::Item, F::Error>;
-}
-
-impl<T, B> ClientExt for Client<T, B>
-    where T: AsyncRead + AsyncWrite + 'static,
-          B: IntoBuf + 'static,
-{
-    fn run<F: Future>(&mut self, f: F) -> Result<F::Item, F::Error> {
-        use futures::future::{self, Future};
-        use futures::future::Either::*;
-
-        let res = future::poll_fn(|| self.poll())
-            .select2(f).wait();
-
-        match res {
-            Ok(A((_, b))) => {
-                // Connection is done...
-                b.wait()
-            }
-            Ok(B((v, _))) => return Ok(v),
-            Err(A((e, _))) => panic!("err: {:?}", e),
-            Err(B((e, _))) => return Err(e),
-        }
-    }
-}
-
-pub mod frames {
-    //! Some useful frames
-
-    pub const SETTINGS: &'static [u8] = &[0, 0, 0, 4, 0, 0, 0, 0, 0];
-    pub const SETTINGS_ACK: &'static [u8] = &[0, 0, 0, 4, 1, 0, 0, 0, 0];
-}
+// This is the frame type that is sent
+pub type SendFrame = h2::frame::Frame<::std::io::Cursor<::bytes::Bytes>>;

tests/support/src/mock.rs

@@ -0,0 +1,305 @@
+use {FutureExt, SendFrame};
+
+use h2::{self, SendError, RecvError};
+use h2::frame::{self, Frame};
+
+use futures::{Future, Stream, Poll};
+use futures::task::{self, Task};
+
+use tokio_io::{AsyncRead, AsyncWrite};
+use tokio_io::io::read_exact;
+
+use std::{cmp, io};
+use std::io::ErrorKind::WouldBlock;
+use std::sync::{Arc, Mutex};
+
+/// A mock I/O
+#[derive(Debug)]
+pub struct Mock {
+    pipe: Pipe,
+}
+
+#[derive(Debug)]
+pub struct Handle {
+    codec: ::Codec<Pipe>,
+}
+
+#[derive(Debug)]
+struct Pipe {
+    inner: Arc<Mutex<Inner>>,
+}
+
+#[derive(Debug)]
+struct Inner {
+    rx: Vec<u8>,
+    rx_task: Option<Task>,
+    tx: Vec<u8>,
+    tx_task: Option<Task>,
+    closed: bool,
+}
+
+const PREFACE: &'static [u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";
+
+/// Create a new mock and handle
+pub fn new() -> (Mock, Handle) {
+    let inner = Arc::new(Mutex::new(Inner {
+        rx: vec![],
+        rx_task: None,
+        tx: vec![],
+        tx_task: None,
+        closed: false,
+    }));
+
+    let mock = Mock {
+        pipe: Pipe {
+            inner: inner.clone(),
+        },
+    };
+
+    let handle = Handle {
+        codec: h2::Codec::new(Pipe { inner }),
+    };
+
+    (mock, handle)
+}
+
+// ===== impl Handle =====
+
+impl Handle {
+    /// Send a frame
+    pub fn send(&mut self, item: SendFrame) -> Result<(), SendError> {
+        // Queue the frame
+        self.codec.buffer(item).unwrap();
+
+        // Flush the frame
+        assert!(self.codec.flush()?.is_ready());
+
+        Ok(())
+    }
+
+    /// Writes the client preface
+    pub fn write_preface(&mut self) {
+        use std::io::Write;
+
+        // Write the connnection preface
+        self.codec.get_mut().write(PREFACE).unwrap();
+    }
+
+    /// Read the client preface
+    pub fn read_preface(self)
+        -> Box<Future<Item = Self, Error = io::Error>>
+    {
+        let buf = vec![0; PREFACE.len()];
+        let ret = read_exact(self, buf)
+            .and_then(|(me, buf)| {
+                assert_eq!(buf, PREFACE);
+                Ok(me)
+            });
+
+        Box::new(ret)
+    }
+
+    /// Perform the H2 handshake
+    pub fn assert_client_handshake(mut self)
+        -> Box<Future<Item = (frame::Settings, Self), Error = h2::Error>>
+    {
+        // Send a settings frame
+        let frame = frame::Settings::default();
+        self.send(frame.into()).unwrap();
+
+        let ret = self.read_preface().unwrap()
+            .and_then(|me| me.into_future().unwrap())
+            .map(|(frame, mut me)| {
+                match frame {
+                    Some(Frame::Settings(settings)) => {
+                        // Send the ACK
+                        let ack = frame::Settings::ack();
+
+                        // TODO: Don't unwrap?
+                        me.send(ack.into()).unwrap();
+
+                        (settings, me)
+                    }
+                    Some(frame) => {
+                        panic!("unexpected frame; frame={:?}", frame);
+                    }
+                    None => {
+                        panic!("unexpected EOF");
+                    }
+                }
+            })
+            .then(|res| {
+                let (settings, me) = res.unwrap();
+
+                me.into_future()
+                    .map_err(|_| unimplemented!())
+                    .map(|(frame, me)| {
+                        let f = assert_settings!(frame.unwrap());
+
+                        // Is ACK
+                        assert!(f.is_ack());
+
+                        (settings, me)
+                    })
+            })
+            ;
+
+        Box::new(ret)
+    }
+}
+
+impl Stream for Handle {
+    type Item = Frame;
+    type Error = RecvError;
+
+    fn poll(&mut self) -> Poll<Option<Self::Item>, RecvError> {
+        self.codec.poll()
+    }
+}
+
+impl io::Read for Handle {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        self.codec.get_mut().read(buf)
+    }
+}
+
+impl AsyncRead for Handle {
+}
+
+impl io::Write for Handle {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.codec.get_mut().write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+impl AsyncWrite for Handle {
+    fn shutdown(&mut self) -> Poll<(), io::Error> {
+        use std::io::Write;
+        try_nb!(self.flush());
+        Ok(().into())
+    }
+}
+
+impl Drop for Handle {
+    fn drop(&mut self) {
+        assert!(self.codec.shutdown().unwrap().is_ready());
+
+        let mut me = self.codec.get_mut().inner.lock().unwrap();
+        me.closed = true;
+
+        if let Some(task) = me.rx_task.take() {
+            task.notify();
+        }
+    }
+}
+
+// ===== impl Mock =====
+
+impl io::Read for Mock {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        assert!(buf.len() > 0, "attempted read with zero length buffer... wut?");
+
+        let mut me = self.pipe.inner.lock().unwrap();
+
+        if me.rx.is_empty() {
+            if me.closed {
+                return Ok(0);
+            }
+
+            me.rx_task = Some(task::current());
+            return Err(WouldBlock.into());
+        }
+
+        let n = cmp::min(buf.len(), me.rx.len());
+        buf[..n].copy_from_slice(&me.rx[..n]);
+        me.rx.drain(..n);
+
+        Ok(n)
+    }
+}
+
+impl AsyncRead for Mock {
+}
+
+impl io::Write for Mock {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        let mut me = self.pipe.inner.lock().unwrap();
+
+        if me.closed {
+            return Err(io::ErrorKind::BrokenPipe.into());
+        }
+
+        me.tx.extend(buf);
+
+        if let Some(task) = me.tx_task.take() {
+            task.notify();
+        }
+
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+impl AsyncWrite for Mock {
+    fn shutdown(&mut self) -> Poll<(), io::Error> {
+        use std::io::Write;
+        try_nb!(self.flush());
+        Ok(().into())
+    }
+}
+
+// ===== impl Pipe =====
+
+impl io::Read for Pipe {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        assert!(buf.len() > 0, "attempted read with zero length buffer... wut?");
+
+        let mut me = self.inner.lock().unwrap();
+
+        if me.tx.is_empty() {
+            me.tx_task = Some(task::current());
+            return Err(WouldBlock.into());
+        }
+
+        let n = cmp::min(buf.len(), me.tx.len());
+        buf[..n].copy_from_slice(&me.tx[..n]);
+        me.tx.drain(..n);
+
+        Ok(n)
+    }
+}
+
+impl AsyncRead for Pipe {
+}
+
+impl io::Write for Pipe {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        let mut me = self.inner.lock().unwrap();
+        me.rx.extend(buf);
+
+        if let Some(task) = me.rx_task.take() {
+            task.notify();
+        }
+
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+impl AsyncWrite for Pipe {
+    fn shutdown(&mut self) -> Poll<(), io::Error> {
+        use std::io::Write;
+        try_nb!(self.flush());
+        Ok(().into())
+    }
+}

tests/support/src/prelude.rs

@@ -0,0 +1,107 @@
+
+// Re-export H2 crate
+pub use super::h2;
+
+pub use self::h2::*;
+pub use self::h2::client::{self, Client};
+pub use self::h2::server::{self, Server};
+
+// Re-export mock
+pub use super::mock;
+
+// Re-export some type defines
+pub use super::{Codec, SendFrame};
+
+// Re-export useful crates
+pub use super::{
+    http,
+    bytes,
+    tokio_io,
+    futures,
+    mock_io,
+    env_logger,
+};
+
+// Re-export primary future types
+pub use self::futures::{
+    Future,
+    Sink,
+    Stream,
+};
+
+// And our Future extensions
+pub use future_ext::{FutureExt, Unwrap};
+
+// Re-export HTTP types
+pub use self::http::{
+    Request,
+    Response,
+    Method,
+    HeaderMap,
+    StatusCode,
+};
+
+pub use self::bytes::{
+    Buf,
+    BufMut,
+    Bytes,
+    BytesMut,
+    IntoBuf,
+};
+
+pub use tokio_io::{AsyncRead, AsyncWrite};
+
+pub use std::time::Duration;
+
+// ===== Everything under here shouldn't be used =====
+// TODO: work on deleting this code
+
+pub use ::futures::future::poll_fn;
+
+pub trait MockH2 {
+    fn handshake(&mut self) -> &mut Self;
+}
+
+impl MockH2 for mock_io::Builder {
+    fn handshake(&mut self) -> &mut Self {
+        self.write(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")
+            // Settings frame
+            .write(frames::SETTINGS)
+            .read(frames::SETTINGS)
+            .read(frames::SETTINGS_ACK)
+    }
+}
+
+pub trait ClientExt {
+    fn run<F: Future>(&mut self, f: F) -> Result<F::Item, F::Error>;
+}
+
+impl<T, B> ClientExt for Client<T, B>
+    where T: AsyncRead + AsyncWrite + 'static,
+          B: IntoBuf + 'static,
+{
+    fn run<F: Future>(&mut self, f: F) -> Result<F::Item, F::Error> {
+        use futures::future::{self, Future};
+        use futures::future::Either::*;
+
+        let res = future::poll_fn(|| self.poll())
+            .select2(f).wait();
+
+        match res {
+            Ok(A((_, b))) => {
+                // Connection is done...
+                b.wait()
+            }
+            Ok(B((v, _))) => return Ok(v),
+            Err(A((e, _))) => panic!("err: {:?}", e),
+            Err(B((e, _))) => return Err(e),
+        }
+    }
+}
+
+pub mod frames {
+    //! Some useful frames
+
+    pub const SETTINGS: &'static [u8] = &[0, 0, 0, 4, 0, 0, 0, 0, 0];
+    pub const SETTINGS_ACK: &'static [u8] = &[0, 0, 0, 4, 1, 0, 0, 0, 0];
+}

tests/support/src/raw.rs

@@ -0,0 +1,46 @@
+// ===== Build a codec from raw bytes =====
+
+#[macro_export]
+macro_rules! raw_codec {
+    (
+        $(
+            $fn:ident => [$($chunk:expr,)+];
+        )*
+    ) => {{
+        let mut b = $crate::mock_io::Builder::new();
+
+        $({
+            let mut chunk = vec![];
+
+            $(
+                $crate::raw::Chunk::push(&$chunk, &mut chunk);
+            )+
+
+            b.$fn(&chunk[..]);
+        })*
+
+        $crate::Codec::new(b.build())
+    }}
+}
+
+pub trait Chunk {
+    fn push(&self, dst: &mut Vec<u8>);
+}
+
+impl Chunk for u8 {
+    fn push(&self, dst: &mut Vec<u8>) {
+        dst.push(*self);
+    }
+}
+
+impl<'a> Chunk for &'a [u8] {
+    fn push(&self, dst: &mut Vec<u8>) {
+        dst.extend(*self)
+    }
+}
+
+impl<'a> Chunk for &'a str {
+    fn push(&self, dst: &mut Vec<u8>) {
+        dst.extend(self.as_bytes())
+    }
+}