feat: add scheduler::flush() to replace sleep(Duration::ZERO) in tests (#4044)

Author: Madoshakalaka

.github/workflows/main-checks.yml

@@ -105,8 +105,8 @@ jobs:
       - name: Run tests - yew
         run: |
           cd packages/yew
-          CHROMEDRIVER=$(which chromedriver) cargo test --features csr,hydration,ssr --target wasm32-unknown-unknown
-          GECKODRIVER=$(which geckodriver) cargo test --features csr,hydration,ssr --target wasm32-unknown-unknown
+          CHROMEDRIVER=$(which chromedriver) cargo test --features csr,hydration,ssr,test --target wasm32-unknown-unknown
+          GECKODRIVER=$(which geckodriver) cargo test --features csr,hydration,ssr,test --target wasm32-unknown-unknown
 
       - name: Run tests - yew-router
         run: |
@@ -254,7 +254,7 @@ jobs:
       - name: Run WASI tests for yew
         run: |
           RUST_LOG=info
-          cargo test --features ssr,hydration --target wasm32-wasip1 -p yew
+          cargo test --features ssr,hydration,test --target wasm32-wasip1 -p yew
 
   example-runnable-tests-on-wasi:
     name: Example Runnable Tests on WASI

packages/yew/Makefile.toml

@@ -1,6 +1,6 @@
 [tasks.native-test]
 command = "cargo"
-args = ["test", "--features", "csr,ssr,hydration"]
+args = ["test", "--features", "csr,ssr,hydration,test"]
 
 [tasks.wasm-test]
 command = "wasm-pack"
@@ -10,7 +10,7 @@ args = [
     "--headless",
     "--",
     "--features",
-    "csr,hydration,ssr",
+    "csr,hydration,ssr,test",
 ]
 
 [tasks.ssr-test]

packages/yew/src/scheduler.rs

@@ -3,6 +3,34 @@
 use std::cell::RefCell;
 use std::collections::BTreeMap;
 use std::rc::Rc;
+#[cfg(any(test, feature = "test"))]
+mod flush_wakers {
+    use std::cell::RefCell;
+    use std::task::Waker;
+
+    thread_local! {
+        static FLUSH_WAKERS: RefCell<Vec<Waker>> = Default::default();
+    }
+
+    #[cfg(all(
+        target_arch = "wasm32",
+        not(target_os = "wasi"),
+        not(feature = "not_browser_env")
+    ))]
+    pub(super) fn register(waker: Waker) {
+        FLUSH_WAKERS.with(|w| {
+            w.borrow_mut().push(waker);
+        });
+    }
+
+    pub(super) fn wake_all() {
+        FLUSH_WAKERS.with(|w| {
+            for waker in w.borrow_mut().drain(..) {
+                waker.wake();
+            }
+        });
+    }
+}
 
 /// Alias for `Rc<RefCell<T>>`
 pub type Shared<T> = Rc<RefCell<T>>;
@@ -207,6 +235,8 @@ pub(crate) fn start_now() {
     LOCK.with(|l| {
         if let Ok(_lock) = l.try_borrow_mut() {
             scheduler_loop();
+            #[cfg(any(test, feature = "test"))]
+            flush_wakers::wake_all();
         }
     });
 }
@@ -232,6 +262,11 @@ mod arch {
         IS_SCHEDULED.store(is, Ordering::Relaxed)
     }
 
+    #[cfg(any(test, feature = "test"))]
+    pub(super) fn is_scheduled() -> bool {
+        check_scheduled()
+    }
+
     /// We delay the start of the scheduler to the end of the micro task queue.
     /// So any messages that needs to be queued can be queued.
     pub(crate) fn start() {
@@ -264,6 +299,51 @@ mod arch {
 
 pub(crate) use arch::*;
 
+/// Flush all pending scheduler work, ensuring all rendering and lifecycle callbacks complete.
+///
+/// On browser WebAssembly targets, the scheduler defers its work to the microtask queue.
+/// This function registers a waker that is notified when `start_now()` finishes draining all
+/// queues, providing proper event-driven render-complete notification without arbitrary sleeps.
+///
+/// On non-browser targets, the scheduler runs synchronously so this simply drains pending work.
+///
+/// Use this in tests after mounting or updating a component to ensure all rendering has
+/// completed before making assertions.
+#[cfg(all(
+    any(test, feature = "test"),
+    target_arch = "wasm32",
+    not(target_os = "wasi"),
+    not(feature = "not_browser_env")
+))]
+pub async fn flush() {
+    std::future::poll_fn(|cx| {
+        start_now();
+
+        if arch::is_scheduled() {
+            flush_wakers::register(cx.waker().clone());
+            std::task::Poll::Pending
+        } else {
+            std::task::Poll::Ready(())
+        }
+    })
+    .await
+}
+
+/// Flush all pending scheduler work, ensuring all rendering and lifecycle callbacks complete.
+///
+/// On non-browser targets, the scheduler runs synchronously so this simply drains pending work.
+#[cfg(all(
+    any(test, feature = "test"),
+    not(all(
+        target_arch = "wasm32",
+        not(target_os = "wasi"),
+        not(feature = "not_browser_env")
+    ))
+))]
+pub async fn flush() {
+    start_now();
+}
+
 impl Scheduler {
     /// Fill vector with tasks to be executed according to Runnable type execution priority
     ///

packages/yew/tests/hydration.rs

@@ -16,7 +16,7 @@ use yew::platform::time::sleep;
 use yew::prelude::*;
 use yew::suspense::{use_future, Suspension, SuspensionResult};
 use yew::virtual_dom::VNode;
-use yew::{component, Renderer, ServerRenderer};
+use yew::{component, scheduler, Renderer, ServerRenderer};
 
 wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
 
@@ -62,12 +62,12 @@ async fn hydration_works() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
         .hydrate();
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     let result = obtain_result_by_id("output");
 
@@ -85,7 +85,7 @@ async fn hydration_works() {
         .unwrap()
         .click();
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     let result = obtain_result_by_id("output");
 
@@ -237,7 +237,7 @@ async fn hydration_with_suspense() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
         .hydrate();
@@ -393,7 +393,7 @@ async fn hydration_with_suspense_not_suspended_at_start() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
         .hydrate();
@@ -524,7 +524,7 @@ async fn hydration_nested_suspense_works() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
         .hydrate();
@@ -661,12 +661,12 @@ async fn hydration_node_ref_works() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
         .hydrate();
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     let result = obtain_result_by_id("output");
     assert_eq!(
@@ -682,7 +682,7 @@ async fn hydration_node_ref_works() {
         .unwrap()
         .click();
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     let result = obtain_result_by_id("output");
     assert_eq!(
@@ -754,16 +754,13 @@ async fn hydration_list_order_works() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
         .hydrate();
 
     // Wait until all suspended components becomes revealed.
-    sleep(Duration::ZERO).await;
-    sleep(Duration::ZERO).await;
-    sleep(Duration::ZERO).await;
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     let result = obtain_result_by_id("output");
     assert_eq!(
@@ -837,13 +834,13 @@ async fn hydration_suspense_no_flickering() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
         .hydrate();
 
     // Wait until all suspended components becomes revealed.
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     let result = obtain_result_by_id("output");
     assert_eq!(
@@ -950,16 +947,13 @@ async fn hydration_order_issue_nested_suspense() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
         .hydrate();
 
     // Wait until all suspended components becomes revealed.
-    sleep(Duration::ZERO).await;
-    sleep(Duration::ZERO).await;
-    sleep(Duration::ZERO).await;
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     let result = obtain_result_by_id("output");
     assert_eq!(
@@ -1180,7 +1174,7 @@ async fn hydration_with_camelcase_svg_elements() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     // Hydrate - this should not panic
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
@@ -1268,15 +1262,12 @@ async fn hydration_suspended_child_does_not_trap_sibling_slot() {
         .unwrap()
         .set_inner_html(&s);
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
         .hydrate();
 
-    sleep(Duration::ZERO).await;
-    sleep(Duration::ZERO).await;
-    sleep(Duration::ZERO).await;
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     let result = obtain_result();
 

packages/yew/tests/mod.rs

@@ -2,12 +2,10 @@
 
 mod common;
 
-use std::time::Duration;
-
 use common::obtain_result;
 use wasm_bindgen_test::*;
-use yew::platform::time::sleep;
 use yew::prelude::*;
+use yew::scheduler;
 
 wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
 
@@ -36,7 +34,7 @@ async fn props_are_passed() {
     )
     .render();
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
     let result = obtain_result();
     assert_eq!(result.as_str(), "done");
 }

packages/yew/tests/suspense.rs

@@ -14,7 +14,7 @@ use yew::platform::spawn_local;
 use yew::platform::time::sleep;
 use yew::prelude::*;
 use yew::suspense::{use_future, use_future_with, Suspension, SuspensionResult};
-use yew::UseStateHandle;
+use yew::{scheduler, UseStateHandle};
 
 wasm_bindgen_test_configure!(run_in_browser);
 
@@ -123,7 +123,7 @@ async fn suspense_works() {
         .unwrap()
         .click();
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     gloo::utils::document()
         .query_selector(".increase")
@@ -553,7 +553,7 @@ async fn effects_not_run_when_suspended() {
         .unwrap()
         .click();
 
-    sleep(Duration::ZERO).await;
+    scheduler::flush().await;
 
     gloo::utils::document()
         .query_selector(".increase")
@@ -563,7 +563,7 @@ async fn effects_not_run_when_suspended() {
         .unwrap()
         .click();
 
-    sleep(Duration::from_millis(0)).await;
+    scheduler::flush().await;
 
     let result = obtain_result();
     assert_eq!(
@@ -788,7 +788,7 @@ async fn test_duplicate_suspension() {
     #[component]
     fn FetchingProvider(props: &ChildrenProps) -> HtmlResult {
         use_future(|| async {
-            sleep(Duration::ZERO).await;
+            scheduler::flush().await;
         })?;
         Ok(html! { <>{props.children.clone()}</> })
     }