Avoid panic when language server is dropped before being initialized in tests

Max Brunsfeld and Nathan Sobo created 3 years ago

Co-authored-by: Nathan Sobo <[e-mail]>

Change summary

crates/language/src/language.rs | 11 +++++++----
crates/lsp/src/lsp.rs           | 10 ++++++++--
crates/project/src/project.rs   |  2 +-
3 files changed, 16 insertions(+), 7 deletions(-)

Detailed changes

crates/language/src/language.rs

@@ -289,10 +289,13 @@ impl LanguageRegistry {
                 let servers_tx = servers_tx.clone();
                 cx.background()
                     .spawn(async move {
-                        fake_server
-                            .receive_notification::<lsp::notification::Initialized>()
-                            .await;
-                        servers_tx.unbounded_send(fake_server).ok();
+                        if fake_server
+                            .try_receive_notification::<lsp::notification::Initialized>()
+                            .await
+                            .is_some()
+                        {
+                            servers_tx.unbounded_send(fake_server).ok();
+                        }
                     })
                     .detach();
                 Ok(server)

crates/lsp/src/lsp.rs

@@ -647,12 +647,18 @@ impl FakeLanguageServer {
     }
 
     pub async fn receive_notification<T: notification::Notification>(&mut self) -> T::Params {
+        self.try_receive_notification::<T>().await.unwrap()
+    }
+
+    pub async fn try_receive_notification<T: notification::Notification>(
+        &mut self,
+    ) -> Option<T::Params> {
         use futures::StreamExt as _;
 
         loop {
-            let (method, params) = self.notifications_rx.next().await.unwrap();
+            let (method, params) = self.notifications_rx.next().await?;
             if &method == T::METHOD {
-                return serde_json::from_str::<T::Params>(&params).unwrap();
+                return Some(serde_json::from_str::<T::Params>(&params).unwrap());
             } else {
                 log::info!("skipping message in fake language server {:?}", params);
             }

crates/project/src/project.rs

@@ -6550,7 +6550,7 @@ mod tests {
         assert!(results.is_empty());
     }
 
-    #[gpui::test]
+    #[gpui::test(iterations = 10)]
     async fn test_definition(cx: &mut gpui::TestAppContext) {
         let mut language = Language::new(
             LanguageConfig {