From f680fca9ec3dae8335f231ebb56061d231ee46ec Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Mon, 12 Mar 2018 14:28:59 -0700
Subject: [PATCH] tests(client): fix flaking conn_reset test

---
 src/client/tests.rs |  9 ++++++++-
 src/mock.rs         | 10 +++++-----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/client/tests.rs b/src/client/tests.rs
index 7d2157dd..daf442f2 100644
--- a/src/client/tests.rs
+++ b/src/client/tests.rs
@@ -1,8 +1,10 @@
 extern crate pretty_env_logger;
 
+use std::time::Duration;
+
 use futures::Async;
 use futures::future::poll_fn;
-use tokio::reactor::Core;
+use tokio::reactor::{Core, Timeout};
 
 use mock::MockConnector;
 use super::*;
@@ -68,6 +70,11 @@ fn conn_reset_after_write() {
             Ok(Async::Ready(()))
         });
         core.run(res1.join(srv1)).expect("res1");
+
+        // run a tiny timeout just to spin the core, so that the pool
+        // can tell the socket is ready again
+        let timeout = Timeout::new(Duration::from_millis(50), &core.handle()).unwrap();
+        core.run(timeout).unwrap();
     }
 
     let res2 = client.get("http://mock.local/a".parse().unwrap());
diff --git a/src/mock.rs b/src/mock.rs
index 2e88b10d..ed340fa0 100644
--- a/src/mock.rs
+++ b/src/mock.rs
@@ -421,10 +421,10 @@ impl Service for MockConnector {
     fn call(&self, uri: Uri) -> Self::Future {
         use futures::future;
         trace!("mock connect: {:?}", uri.as_ref());
-        let mock = self.mocks.borrow_mut()
-            .get_mut(uri.as_ref())
-            .expect(&format!("unknown mocks uri: {:?}", uri.as_ref()))
-            .remove(0);
-        future::ok(mock)
+        let mut mocks = self.mocks.borrow_mut();
+        let mocks = mocks.get_mut(uri.as_ref())
+            .expect(&format!("unknown mocks uri: {:?}", uri.as_ref()));
+        assert!(!mocks.is_empty(), "no additional mocks for {:?}", uri.as_ref());
+        future::ok(mocks.remove(0))
     }
 }