Fix panic because of unwrapping when resuming suspension (#2902)

* Replace unwrap with if let to prevent panic

* Add resume after unmount test

* Change BaseSuspense methods to associated functions

* Fix rustfmt issues

packages/yew/src/html/component/lifecycle.rs

@@ -303,9 +303,7 @@ impl ComponentState {
             let comp_scope = self.inner.any_scope();
 
             let suspense_scope = comp_scope.find_parent_scope::<BaseSuspense>().unwrap();
-            let suspense = suspense_scope.get_component().unwrap();
-
-            suspense.resume(m);
+            BaseSuspense::resume(&suspense_scope, m);
         }
     }
 }
@@ -462,7 +460,6 @@ impl ComponentState {
             let suspense_scope = comp_scope
                 .find_parent_scope::<BaseSuspense>()
                 .expect("To suspend rendering, a <Suspense /> component is required.");
-            let suspense = suspense_scope.get_component().unwrap();
 
             let comp_id = self.comp_id;
             let shared_state = shared_state.clone();
@@ -479,12 +476,12 @@ impl ComponentState {
             if let Some(ref last_suspension) = self.suspension {
                 if &suspension != last_suspension {
                     // We remove previous suspension from the suspense.
-                    suspense.resume(last_suspension.clone());
+                    BaseSuspense::resume(&suspense_scope, last_suspension.clone());
                 }
             }
             self.suspension = Some(suspension.clone());
 
-            suspense.suspend(suspension);
+            BaseSuspense::suspend(&suspense_scope, suspension);
         }
     }
 

packages/yew/src/suspense/component.rs

@@ -36,7 +36,6 @@ mod feat_csr_ssr {
 
     #[derive(Debug)]
     pub(crate) struct BaseSuspense {
-        link: Scope<Self>,
         suspensions: Vec<Suspension>,
         #[cfg(feature = "hydration")]
         hydration_handle: Option<SuspensionHandle>,
@@ -46,7 +45,7 @@ mod feat_csr_ssr {
         type Message = BaseSuspenseMsg;
         type Properties = BaseSuspenseProps;
 
-        fn create(ctx: &Context<Self>) -> Self {
+        fn create(_ctx: &Context<Self>) -> Self {
             #[cfg(not(feature = "hydration"))]
             let suspensions = Vec::new();
 
@@ -56,9 +55,9 @@ mod feat_csr_ssr {
                 use crate::callback::Callback;
                 use crate::html::RenderMode;
 
-                match ctx.creation_mode() {
+                match _ctx.creation_mode() {
                     RenderMode::Hydration => {
-                        let link = ctx.link().clone();
+                        let link = _ctx.link().clone();
                         let (s, handle) = Suspension::new();
                         s.listen(Callback::from(move |s| {
                             link.send_message(BaseSuspenseMsg::Resume(s));
@@ -70,7 +69,6 @@ mod feat_csr_ssr {
             };
 
             Self {
-                link: ctx.link().clone(),
                 suspensions,
                 #[cfg(feature = "hydration")]
                 hydration_handle,
@@ -133,12 +131,12 @@ mod feat_csr_ssr {
     }
 
     impl BaseSuspense {
-        pub(crate) fn suspend(&self, s: Suspension) {
-            self.link.send_message(BaseSuspenseMsg::Suspend(s));
+        pub(crate) fn suspend(scope: &Scope<Self>, s: Suspension) {
+            scope.send_message(BaseSuspenseMsg::Suspend(s));
         }
 
-        pub(crate) fn resume(&self, s: Suspension) {
-            self.link.send_message(BaseSuspenseMsg::Resume(s));
+        pub(crate) fn resume(scope: &Scope<Self>, s: Suspension) {
+            scope.send_message(BaseSuspenseMsg::Resume(s));
         }
     }
 

packages/yew/tests/suspense.rs

@@ -14,6 +14,7 @@ use yew::platform::spawn_local;
 use yew::platform::time::sleep;
 use yew::prelude::*;
 use yew::suspense::{use_future, use_future_with_deps, Suspension, SuspensionResult};
+use yew::UseStateHandle;
 
 wasm_bindgen_test_configure!(run_in_browser);
 
@@ -734,3 +735,54 @@ async fn test_suspend_forever() {
     let result = obtain_result();
     assert_eq!(result.as_str(), r#"OK"#);
 }
+
+#[wasm_bindgen_test]
+async fn resume_after_unmount() {
+    #[derive(Clone, Properties, PartialEq)]
+    struct ContentProps {
+        state: UseStateHandle<bool>,
+    }
+
+    #[function_component(Content)]
+    fn content(ContentProps { state }: &ContentProps) -> HtmlResult {
+        let state = state.clone();
+        let _sleep_handle = use_future(|| async move {
+            sleep(Duration::from_millis(50)).await;
+            state.set(false);
+            sleep(Duration::from_millis(50)).await;
+        })?;
+
+        Ok(html! {
+            <div>{"Content"}</div>
+        })
+    }
+
+    #[function_component(App)]
+    fn app() -> Html {
+        let fallback = html! {<div>{"wait..."}</div>};
+        let state = use_state(|| true);
+
+        html! {
+            <div id="result">
+            if *state {
+                <Suspense {fallback}>
+                    <Content {state} />
+                </Suspense>
+            } else {
+                <div>{"Content replacement"}</div>
+            }
+            </div>
+        }
+    }
+
+    yew::Renderer::<App>::with_root(gloo::utils::document().get_element_by_id("output").unwrap())
+        .render();
+
+    sleep(Duration::from_millis(25)).await;
+    let result = obtain_result();
+    assert_eq!(result.as_str(), "<div>wait...</div>");
+
+    sleep(Duration::from_millis(50)).await;
+    let result = obtain_result();
+    assert_eq!(result.as_str(), "<div>Content replacement</div>");
+}