From 90df6e387901e85f7f7d0cbadcdadc72d0db1afe Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Fri, 4 Aug 2017 20:54:49 -0700
Subject: [PATCH] Try to clean up test

---
 tests/stream_states.rs |  7 +++----
 tests/support/mod.rs   | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/tests/stream_states.rs b/tests/stream_states.rs
index c875297..1a4afe0 100644
--- a/tests/stream_states.rs
+++ b/tests/stream_states.rs
@@ -32,13 +32,12 @@ fn send_recv_headers_only() {
         .body(()).unwrap();
 
     info!("sending request");
-    let stream = h2.request(request, true).unwrap();
+    let mut stream = h2.request(request, true).unwrap();
 
-    // Wait
-    h2.wait().ok().unwrap();
+    let resp = h2.run(poll_fn(|| stream.poll_response())).unwrap();
 
     // Try to get response
-    println!("GOT: {:?}", stream.wait().ok().unwrap());
+    println!("GOT: {:?}", resp);
 
     // let resp = stream.select2(h2).wait().ok().unwrap();
 
diff --git a/tests/support/mod.rs b/tests/support/mod.rs
index 0d2ebf6..47869d8 100644
--- a/tests/support/mod.rs
+++ b/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