From e8ea2441a4a6faedb9f83a27763e2d6e9f4fd7c9 Mon Sep 17 00:00:00 2001
From: fakeshadow <24548779@qq.com>
Date: Thu, 25 Mar 2021 22:43:31 +0800
Subject: [PATCH] fix test

---
 actix-http/src/client/connection.rs | 48 ++++++++++++++++++++++++-----
 1 file changed, 41 insertions(+), 7 deletions(-)

diff --git a/actix-http/src/client/connection.rs b/actix-http/src/client/connection.rs
index a02f88e25..0e3e97f3f 100644
--- a/actix-http/src/client/connection.rs
+++ b/actix-http/src/client/connection.rs
@@ -404,9 +404,18 @@ where
 
 #[cfg(test)]
 mod test {
-    use std::net;
+    use std::{
+        future::Future,
+        net,
+        pin::Pin,
+        task::{Context, Poll},
+        time::{Duration, Instant},
+    };
 
-    use actix_rt::net::TcpStream;
+    use actix_rt::{
+        net::TcpStream,
+        time::{interval, Interval},
+    };
 
     use super::*;
 
@@ -430,11 +439,36 @@ mod test {
 
         drop(conn);
 
-        actix_rt::task::yield_now().await;
+        struct DropCheck {
+            sender: h2::client::SendRequest<Bytes>,
+            interval: Interval,
+            start_from: Instant,
+        }
 
-        match sender.ready().await {
-            Ok(_) => panic!("connection should be gone and can not be ready"),
-            Err(e) => assert!(e.is_io()),
-        };
+        impl Future for DropCheck {
+            type Output = ();
+
+            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+                let this = self.get_mut();
+                match futures_core::ready!(this.sender.poll_ready(cx)) {
+                    Ok(()) => {
+                        if this.start_from.elapsed() > Duration::from_secs(10) {
+                            panic!("connection should be gone and can not be ready");
+                        } else {
+                            let _ = this.interval.poll_tick(cx);
+                            Poll::Pending
+                        }
+                    }
+                    Err(_) => Poll::Ready(()),
+                }
+            }
+        }
+
+        DropCheck {
+            sender,
+            interval: interval(Duration::from_millis(100)),
+            start_from: Instant::now(),
+        }
+        .await;
     }
 }