Fix tests.

yewstack · Aug 5, 2022 · 892b759 · 892b759
1 parent fb2603e
commit 892b759
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 52 deletions.
diff --git a/packages/yew/src/platform/pinned/mpsc.rs b/packages/yew/src/platform/pinned/mpsc.rs
@@ -255,6 +255,7 @@ mod tests {
 
     use futures::sink::SinkExt;
     use futures::stream::StreamExt;
+    use tokio::task::LocalSet;
     use tokio::test;
 
     use super::*;
@@ -263,22 +264,28 @@ mod tests {
 
     #[test]
     async fn mpsc_works() {
-        let (tx, mut rx) = unbounded::<usize>();
+        let local_set = LocalSet::new();
 
-        spawn_local(async move {
-            for i in 0..10 {
-                (&tx).send(i).await.expect("failed to send.");
-                sleep(Duration::from_millis(1)).await;
-            }
-        });
+        local_set
+            .run_until(async {
+                let (tx, mut rx) = unbounded::<usize>();
 
-        for i in 0..10 {
-            let received = rx.next().await.expect("failed to receive");
+                spawn_local(async move {
+                    for i in 0..10 {
+                        (&tx).send(i).await.expect("failed to send.");
+                        sleep(Duration::from_millis(1)).await;
+                    }
+                });
 
-            assert_eq!(i, received);
-        }
+                for i in 0..10 {
+                    let received = rx.next().await.expect("failed to receive");
 
-        assert_eq!(rx.next().await, None);
+                    assert_eq!(i, received);
+                }
+
+                assert_eq!(rx.next().await, None);
+            })
+            .await;
     }
 
     #[test]
@@ -291,37 +298,43 @@ mod tests {
 
     #[test]
     async fn mpsc_multi_sender() {
-        let (tx, mut rx) = unbounded::<usize>();
+        let local_set = LocalSet::new();
 
-        spawn_local(async move {
-            let tx2 = tx.clone();
+        local_set
+            .run_until(async {
+                let (tx, mut rx) = unbounded::<usize>();
 
-            for i in 0..10 {
-                if i % 2 == 0 {
-                    (&tx).send(i).await.expect("failed to send.");
-                } else {
-                    (&tx2).send(i).await.expect("failed to send.");
-                }
+                spawn_local(async move {
+                    let tx2 = tx.clone();
 
-                sleep(Duration::from_millis(1)).await;
-            }
+                    for i in 0..10 {
+                        if i % 2 == 0 {
+                            (&tx).send(i).await.expect("failed to send.");
+                        } else {
+                            (&tx2).send(i).await.expect("failed to send.");
+                        }
 
-            drop(tx2);
+                        sleep(Duration::from_millis(1)).await;
+                    }
 
-            for i in 10..20 {
-                (&tx).send(i).await.expect("failed to send.");
+                    drop(tx2);
 
-                sleep(Duration::from_millis(1)).await;
-            }
-        });
+                    for i in 10..20 {
+                        (&tx).send(i).await.expect("failed to send.");
 
-        for i in 0..20 {
-            let received = rx.next().await.expect("failed to receive");
+                        sleep(Duration::from_millis(1)).await;
+                    }
+                });
 
-            assert_eq!(i, received);
-        }
+                for i in 0..20 {
+                    let received = rx.next().await.expect("failed to receive");
 
-        assert_eq!(rx.next().await, None);
+                    assert_eq!(i, received);
+                }
+
+                assert_eq!(rx.next().await, None);
+            })
+            .await;
     }
 
     #[test]

diff --git a/packages/yew/src/platform/pinned/oneshot.rs b/packages/yew/src/platform/pinned/oneshot.rs
@@ -147,6 +147,7 @@ mod tests {
     use std::time::Duration;
 
     use tokio::sync::Barrier;
+    use tokio::task::LocalSet;
     use tokio::test;
 
     use super::*;
@@ -164,35 +165,47 @@ mod tests {
 
     #[test]
     async fn oneshot_drops_sender() {
-        let (tx, rx) = channel::<usize>();
+        let local_set = LocalSet::new();
 
-        spawn_local(async move {
-            sleep(Duration::from_millis(1)).await;
+        local_set
+            .run_until(async {
+                let (tx, rx) = channel::<usize>();
 
-            drop(tx);
-        });
-        rx.await.expect_err("successful to receive.");
+                spawn_local(async move {
+                    sleep(Duration::from_millis(1)).await;
+
+                    drop(tx);
+                });
+                rx.await.expect_err("successful to receive.");
+            })
+            .await;
     }
 
     #[test]
     async fn oneshot_drops_receiver() {
-        let (tx, rx) = channel::<usize>();
+        let local_set = LocalSet::new();
 
-        let bar = Arc::new(Barrier::new(2));
+        local_set
+            .run_until(async {
+                let (tx, rx) = channel::<usize>();
 
-        {
-            let bar = bar.clone();
-            spawn_local(async move {
-                sleep(Duration::from_millis(1)).await;
+                let bar = Arc::new(Barrier::new(2));
 
-                drop(rx);
+                {
+                    let bar = bar.clone();
+                    spawn_local(async move {
+                        sleep(Duration::from_millis(1)).await;
 
-                bar.wait().await;
-            });
-        }
+                        drop(rx);
 
-        bar.wait().await;
+                        bar.wait().await;
+                    });
+                }
+
+                bar.wait().await;
 
-        tx.send(0).expect_err("successful to send.");
+                tx.send(0).expect_err("successful to send.");
+            })
+            .await;
     }
 }