From 36be5eeec8b9db00e0bac85ce32a1f563144a728 Mon Sep 17 00:00:00 2001
From: j-mendez <jeff@a11ywatch.com>
Date: Wed, 17 Jan 2024 08:41:54 -0500
Subject: [PATCH] feat(page): add wait_for_network_idle_with_timeout

---
 src/page.rs | 60 ++++++++++++++++++++++++++++++++---------------------
 1 file changed, 36 insertions(+), 24 deletions(-)

diff --git a/src/page.rs b/src/page.rs
index 90cd7aaa..75aabf90 100644
--- a/src/page.rs
+++ b/src/page.rs
@@ -308,19 +308,15 @@ impl Page {
     pub async fn wait_for_network_idle(&self) -> Result<&Self> {
         let mut events = self.event_listener::<chromiumoxide_cdp::cdp::browser_protocol::network::EventLoadingFinished>().await?;
 
-        if let Err(_) = tokio::time::timeout(tokio::time::Duration::from_secs(30), async move {
-            loop {
-                let sleep = tokio::time::sleep(tokio::time::Duration::from_millis(500));
-                tokio::pin!(sleep);
-                tokio::select! {
-                    _ = &mut sleep => break,
-                    _ = events.next() => (),
-                    else => break,
-                }
+        loop {
+            let sleep = tokio::time::sleep(tokio::time::Duration::from_millis(500));
+            tokio::pin!(sleep);
+            tokio::select! {
+                _ = &mut sleep => break,
+                _ = events.next() => (),
+                else => break,
             }
-        })
-        .await
-        {}
+        }
 
         Ok(self)
     }
@@ -331,22 +327,38 @@ impl Page {
         use futures::{future::FutureExt, pin_mut, select};
         let mut events = self.event_listener::<chromiumoxide_cdp::cdp::browser_protocol::network::EventLoadingFinished>().await?;
 
-        async_std::io::timeout(std::time::Duration::from_secs(30), async {
-            loop {
-                let t1 = async_std::task::sleep(std::time::Duration::from_millis(500)).fuse();
-                let t2 = events.next().fuse();
+        loop {
+            let t1 = async_std::task::sleep(std::time::Duration::from_millis(500)).fuse();
+            let t2 = events.next().fuse();
 
-                pin_mut!(t1, t2);
+            pin_mut!(t1, t2);
 
-                select! {
-                    () = t1 => break,
-                    _ = t2 => (),
-                }
+            select! {
+                () = t1 => break,
+                _ = t2 => (),
             }
-            Ok(())
-        })
-        .await?;
+        }
+
+        Ok(self)
+    }
 
+    /// Wait for the network to be idle for 500ms with a duration timeout
+    #[cfg(feature = "tokio-runtime")]
+    pub async fn wait_for_network_idle_with_timeout(
+        &self,
+        timeout: core::time::Duration,
+    ) -> Result<&Self> {
+        if let Err(_) = tokio::time::timeout(timeout, self.wait_for_network_idle()).await {}
+        Ok(self)
+    }
+
+    /// Wait for the network to be idle for 500ms with a duration timeout
+    #[cfg(feature = "async-std-runtime")]
+    pub async fn wait_for_network_idle_with_timeout(
+        &self,
+        timeout: core::time::Duration,
+    ) -> Result<&Self> {
+        let _ = (async_std::future::timeout(timeout, self.wait_for_network_idle()).await).is_err();
         Ok(self)
     }