Try to clean up test

tests/support/mod.rs

@@ -3,6 +3,7 @@
 pub extern crate bytes;
 pub extern crate h2;
 pub extern crate http;
+pub extern crate tokio_io;
 pub extern crate futures;
 pub extern crate mock_io;
 pub extern crate env_logger;
@@ -12,6 +13,7 @@ pub use self::futures::{
     Sink,
     Stream,
 };
+pub use self::futures::future::poll_fn;
 
 pub use self::http::{
     request,
@@ -30,8 +32,11 @@ pub use self::bytes::{
     BufMut,
     Bytes,
     BytesMut,
+    IntoBuf,
 };
 
+use tokio_io::{AsyncRead, AsyncWrite};
+
 pub trait MockH2 {
     fn handshake(&mut self) -> &mut Self;
 }
@@ -46,6 +51,33 @@ impl MockH2 for mock_io::Builder {
     }
 }
 
+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