diff --git a/tests/support/src/mock.rs b/tests/support/src/mock.rs index cc53f76..b4e4eb5 100644 --- a/tests/support/src/mock.rs +++ b/tests/support/src/mock.rs @@ -3,13 +3,13 @@ use {FutureExt, SendFrame}; use h2::{self, SendError, RecvError}; use h2::frame::{self, Frame}; -use futures::{Future, Stream, Poll}; +use futures::{Async, 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::{cmp, fmt, io}; use std::io::ErrorKind::WouldBlock; use std::sync::{Arc, Mutex}; @@ -303,3 +303,118 @@ impl AsyncWrite for Pipe { Ok(().into()) } } + +pub trait HandleFutureExt { + fn recv_settings(self) -> RecvFrame, Handle), Error=()>>> + where Self: Sized + 'static, + Self: Future, + Self::Error: fmt::Debug, + { + let map = self.map(|(settings, handle)| (Some(settings.into()), handle)).unwrap(); + + let boxed: Box, Handle), Error=()>> = Box::new(map); + RecvFrame { + inner: boxed, + frame: frame::Settings::default().into(), + } + } + + fn recv_frame(self, frame: T) -> RecvFrame<::Future> + where Self: IntoRecvFrame + Sized, + T: Into, + { + self.into_recv_frame(frame.into()) + } + + fn send_frame(self, frame: T) -> SendFrameFut + where Self: Sized, + T: Into, + { + SendFrameFut { + inner: self, + frame: Some(frame.into()), + } + } +} + +pub struct RecvFrame { + inner: T, + frame: Frame, +} + +impl Future for RecvFrame + where T: Future, Handle)>, + T::Error: fmt::Debug, +{ + type Item = Handle; + type Error = (); + + fn poll(&mut self) -> Poll { + let (frame, handle) = match self.inner.poll().unwrap() { + Async::Ready((frame, handle)) => (frame, handle), + Async::NotReady => return Ok(Async::NotReady), + }; + assert_eq!(frame.unwrap(), self.frame); + Ok(Async::Ready(handle)) + } +} + +pub struct SendFrameFut { + inner: T, + frame: Option, +} + +impl Future for SendFrameFut + where T: Future, + T::Error: fmt::Debug, +{ + type Item = Handle; + type Error = (); + + fn poll(&mut self) -> Poll { + let mut handle = match self.inner.poll().unwrap() { + Async::Ready(handle) => handle, + Async::NotReady => return Ok(Async::NotReady), + }; + handle.send(self.frame.take().unwrap()).unwrap(); + Ok(Async::Ready(handle)) + } +} + +impl HandleFutureExt for T + where T: Future + 'static, +{ +} + +pub trait IntoRecvFrame { + type Future: Future; + fn into_recv_frame(self, frame: Frame) -> RecvFrame; +} + +impl IntoRecvFrame for Handle { + type Future = ::futures::stream::StreamFuture; + + fn into_recv_frame(self, frame: Frame) -> RecvFrame { + RecvFrame { + inner: self.into_future(), + frame: frame, + } + } +} + +impl IntoRecvFrame for T + where T: Future + 'static, + T::Error: fmt::Debug, +{ + type Future = Box, Handle), Error=()>>; + + fn into_recv_frame(self, frame: Frame) -> RecvFrame { + let into_fut = Box::new(self.unwrap() + .and_then(|handle| handle.into_future().unwrap()) + ); + RecvFrame { + inner: into_fut, + frame: frame, + } + } +} diff --git a/tests/support/src/prelude.rs b/tests/support/src/prelude.rs index e3cad03..12bd876 100644 --- a/tests/support/src/prelude.rs +++ b/tests/support/src/prelude.rs @@ -7,7 +7,7 @@ pub use self::h2::client::{self, Client}; pub use self::h2::server::{self, Server}; // Re-export mock -pub use super::mock; +pub use super::mock::{self, HandleFutureExt}; // Re-export frames helpers pub use super::frames;