From c70f3e4ffe210901b1d375db29bd7cce6d3e4cc5 Mon Sep 17 00:00:00 2001
From: Piotr Magiera <piotrmagiera150@gmail.com>
Date: Wed, 9 Oct 2024 13:10:04 +0200
Subject: [PATCH] es

---
 crates/cairo-lang-language-server/src/lib.rs | 61 ++++++++++++--------
 1 file changed, 38 insertions(+), 23 deletions(-)

diff --git a/crates/cairo-lang-language-server/src/lib.rs b/crates/cairo-lang-language-server/src/lib.rs
index d77817c53bb..b46fbbc331c 100644
--- a/crates/cairo-lang-language-server/src/lib.rs
+++ b/crates/cairo-lang-language-server/src/lib.rs
@@ -158,9 +158,9 @@ pub fn start_with_tricks(tricks: Tricks) -> ExitCode {
 
 /// Special function to run the language server in end-to-end tests.
 #[cfg(feature = "testing")]
-pub fn build_service_for_e2e_tests() -> (Box<dyn FnOnce() -> Backend + Send>, lsp_server::Connection)
-{
-    Backend::new_for_testing(Default::default())
+pub fn build_service_for_e2e_tests()
+-> (Box<dyn FnOnce() -> BackendForTesting + Send>, lsp_server::Connection) {
+    BackendForTesting::new_for_testing(Default::default())
 }
 
 /// Initialize logging infrastructure for the language server.
@@ -246,36 +246,52 @@ fn ensure_exists_in_db<'a>(
     new_db.set_file_overrides(Arc::new(new_overrides));
 }
 
-#[cfg(not(feature = "testing"))]
 struct Backend {
     connection: Connection,
     state: State,
 }
 
 #[cfg(feature = "testing")]
-pub struct Backend {
-    connection: Connection,
-    state: State,
-}
-
-impl Backend {
-    fn new(tricks: Tricks) -> Result<Self> {
-        let connection_initializer = ConnectionInitializer::stdio();
-
-        Self::new_inner(tricks, connection_initializer)
-    }
+pub struct BackendForTesting(Backend);
 
-    #[cfg(feature = "testing")]
+#[cfg(feature = "testing")]
+impl BackendForTesting {
     fn new_for_testing(
         tricks: Tricks,
-    ) -> (Box<dyn FnOnce() -> Self + Send>, lsp_server::Connection) {
+    ) -> (Box<dyn FnOnce() -> BackendForTesting + Send>, lsp_server::Connection) {
         let (connection_initializer, client) = ConnectionInitializer::memory();
 
-        let init = Box::new(|| Self::new_inner(tricks, connection_initializer).unwrap());
+        let init = Box::new(|| {
+            BackendForTesting(Backend::new_inner(tricks, connection_initializer).unwrap())
+        });
 
         (init, client)
     }
 
+    pub fn run_for_tests(self) -> Result<JoinHandle<Result<()>>> {
+        let Backend { mut state, connection } = self.0;
+
+        event_loop_thread(move || {
+            let scheduler = Backend::initial_setup(&mut state, &connection);
+
+            let result = Backend::event_loop(&connection, scheduler);
+
+            if let Err(err) = connection.close() {
+                error!("failed to close connection to the language server: {err:?}");
+            }
+
+            result
+        })
+    }
+}
+
+impl Backend {
+    fn new(tricks: Tricks) -> Result<Self> {
+        let connection_initializer = ConnectionInitializer::stdio();
+
+        Self::new_inner(tricks, connection_initializer)
+    }
+
     fn new_inner(tricks: Tricks, connection_initializer: ConnectionInitializer) -> Result<Self> {
         let (id, init_params) = connection_initializer.initialize_start()?;
 
@@ -316,11 +332,6 @@ impl Backend {
         })
     }
 
-    #[cfg(feature = "testing")]
-    pub fn run_for_tests(self) -> Result<JoinHandle<Result<()>>> {
-        self.run()
-    }
-
     fn initial_setup<'a>(state: &'a mut State, connection: &'_ Connection) -> Scheduler<'a> {
         let four = NonZeroUsize::new(4).unwrap();
         let worker_threads = std::thread::available_parallelism().unwrap_or(four).max(four);
@@ -329,6 +340,10 @@ impl Backend {
 
         let mut scheduler = Scheduler::new(state, worker_threads, connection.make_sender());
 
+        // The dynamic registration and config-reload happen here instead of as a reaction to
+        // `initialized` notification as it is more convenient implementation-wise.
+        // Ideally, we would wait for responses to these requests here instead of in the event loop,
+        // but we did not find a suitable way to do that.
         if let Err(error) =
             Self::register_dynamic_capabilities(&mut scheduler, dynamic_registrations)
         {