fix and test spawn_forever (#2216)

packages/core/tests/task.rs

@@ -49,6 +49,47 @@ async fn running_async() {
     );
 }
 
+#[tokio::test]
+async fn spawn_forever_persists() {
+    use std::sync::atomic::Ordering;
+    static POLL_COUNT: AtomicUsize = AtomicUsize::new(0);
+
+    fn app() -> Element {
+        if generation() > 0 {
+            rsx!(div {})
+        } else {
+            needs_update();
+            rsx!(Child {})
+        }
+    }
+
+    #[component]
+    fn Child() -> Element {
+        spawn_forever(async move {
+            loop {
+                POLL_COUNT.fetch_add(1, Ordering::Relaxed);
+                tokio::time::sleep(Duration::from_millis(50)).await;
+            }
+        });
+
+        rsx!(div {})
+    }
+
+    let mut dom = VirtualDom::new(app);
+
+    dom.rebuild(&mut dioxus_core::NoOpMutations);
+    dom.render_immediate(&mut dioxus_core::NoOpMutations);
+
+    tokio::select! {
+        _ = dom.wait_for_work() => {}
+        _ = tokio::time::sleep(Duration::from_millis(500)) => {}
+    };
+
+    // By the time the tasks are finished, we should've accumulated ticks from two tasks
+    // Be warned that by setting the delay to too short, tokio might not schedule in the tasks
+    assert_eq!(POLL_COUNT.load(Ordering::Relaxed), 10);
+}
+
 /// Prove that yield_now doesn't cause a deadlock
 #[tokio::test]
 async fn yield_now_works() {