Make testing more reliable and realistic

crates/air/src/commands/lsp.rs

@@ -4,7 +4,7 @@ use crate::ExitStatus;
 #[tokio::main]
 pub(crate) async fn lsp(_command: LspCommand) -> anyhow::Result<ExitStatus> {
     // Returns after shutdown
-    lsp::start_lsp(tokio::io::stdin(), tokio::io::stdout()).await;
+    lsp::start_server(tokio::io::stdin(), tokio::io::stdout()).await;
 
     Ok(ExitStatus::Success)
 }

crates/lsp/src/handlers_format.rs

@@ -40,51 +40,3 @@ pub(crate) fn document_formatting(
     let edits = to_proto::replace_all_edit(&doc.line_index, &doc.contents, &output)?;
     Ok(Some(edits))
 }
-
-#[cfg(test)]
-mod tests {
-    use crate::{
-        documents::Document, tower_lsp::init_test_client, tower_lsp_test_client::TestClientExt,
-    };
-
-    #[tests_macros::lsp_test]
-    async fn test_format() {
-        let mut client = init_test_client().await;
-
-        #[rustfmt::skip]
-        let doc = Document::doodle(
-"
-1
-2+2
-3 + 3 +
-3",
-        );
-
-        let formatted = client.format_document(&doc).await;
-        insta::assert_snapshot!(formatted);
-
-        client
-    }
-
-    // https://github.com/posit-dev/air/issues/61
-    #[tests_macros::lsp_test]
-    async fn test_format_minimal_diff() {
-        let mut client = init_test_client().await;
-
-        #[rustfmt::skip]
-        let doc = Document::doodle(
-"1
-2+2
-3
-",
-        );
-
-        let edits = client.format_document_edits(&doc).await.unwrap();
-        assert!(edits.len() == 1);
-
-        let edit = &edits[0];
-        assert_eq!(edit.new_text, " + ");
-
-        client
-    }
-}

crates/lsp/src/lib.rs

@@ -1,7 +1,9 @@
 // TODO: Remove this
 #![allow(dead_code)]
 
-pub use tower_lsp::start_lsp;
+use std::sync::OnceLock;
+
+pub use tower_lsp::start_server;
 
 pub mod config;
 pub mod documents;
@@ -17,8 +19,11 @@ pub mod state;
 pub mod to_proto;
 pub mod tower_lsp;
 
-#[cfg(test)]
-pub mod tower_lsp_test_client;
+/// Are we in a test LSP server?
+///
+/// `start_test_server()` sets this to `true`, which disables logging to the LSP client
+/// during tests, making it easier to track sent/received requests.
+pub static TESTING: OnceLock<bool> = OnceLock::new();
 
 // These send LSP messages in a non-async and non-blocking way.
 // The LOG level is not timestamped so we're not using it.

crates/lsp/src/main_loop.rs

@@ -8,6 +8,7 @@
 use std::collections::HashMap;
 use std::future;
 use std::pin::Pin;
+use std::sync::OnceLock;
 
 use anyhow::anyhow;
 use biome_lsp_converters::PositionEncoding;
@@ -31,16 +32,15 @@ use crate::tower_lsp::LspMessage;
 use crate::tower_lsp::LspNotification;
 use crate::tower_lsp::LspRequest;
 use crate::tower_lsp::LspResponse;
+use crate::TESTING;
 
 pub(crate) type TokioUnboundedSender<T> = tokio::sync::mpsc::UnboundedSender<T>;
 pub(crate) type TokioUnboundedReceiver<T> = tokio::sync::mpsc::UnboundedReceiver<T>;
 
 // The global instance of the auxiliary event channel, used for sending log
 // messages or spawning threads from free functions. Since this is an unbounded
-// channel, sending a log message is not async nor blocking. Tokio senders are
-// Send and Sync so this global variable can be safely shared across threads.
-static mut AUXILIARY_EVENT_TX: std::cell::OnceCell<TokioUnboundedSender<AuxiliaryEvent>> =
-    std::cell::OnceCell::new();
+// channel, sending a log message is not async nor blocking.
+static AUXILIARY_EVENT_TX: OnceLock<TokioUnboundedSender<AuxiliaryEvent>> = OnceLock::new();
 
 // This is the syntax for trait aliases until an official one is stabilised.
 // This alias is for the future of a `JoinHandle<anyhow::Result<T>>`
@@ -378,16 +378,9 @@ impl AuxiliaryState {
         // Set global instance of this channel. This is used for interacting
         // with the auxiliary loop (logging messages or spawning a task) from
         // free functions.
-        unsafe {
-            #[allow(static_mut_refs)]
-            if let Some(val) = AUXILIARY_EVENT_TX.get_mut() {
-                // Reset channel if already set. Happens e.g. on reconnection after a refresh.
-                *val = auxiliary_event_tx;
-            } else {
-                // Set channel for the first time
-                AUXILIARY_EVENT_TX.set(auxiliary_event_tx).unwrap();
-            }
-        }
+        AUXILIARY_EVENT_TX
+            .set(auxiliary_event_tx)
+            .expect("Auxiliary event channel can't be set more than once.");
 
         // List of pending tasks for which we manage the lifecycle (mainly relay
         // errors and panics)
@@ -453,12 +446,10 @@ impl AuxiliaryState {
 }
 
 fn auxiliary_tx() -> &'static TokioUnboundedSender<AuxiliaryEvent> {
-    // If we get here that means the LSP was initialised at least once. The
-    // channel might be closed if the LSP was dropped, but it should exist.
-    unsafe {
-        #[allow(static_mut_refs)]
-        AUXILIARY_EVENT_TX.get().unwrap()
-    }
+    // If we get here that means the LSP was initialised in `AuxiliaryState::new()`.
+    // The channel might be closed if the LSP was dropped, but it should exist
+    // (and in that case we expect the process to exit shortly afterwards anyways).
+    AUXILIARY_EVENT_TX.get().unwrap()
 }
 
 fn send_auxiliary(event: AuxiliaryEvent) {
@@ -471,8 +462,9 @@ fn send_auxiliary(event: AuxiliaryEvent) {
 /// Send a message to the LSP client. This is non-blocking and treated on a
 /// latency-sensitive task.
 pub(crate) fn log(level: lsp_types::MessageType, message: String) {
-    // We're not connected to an LSP client when running unit tests
-    if cfg!(test) {
+    // We don't want to send logs to the client when running integration tests,
+    // as they interfere with our ability to track sent/received requests.
+    if *TESTING.get().unwrap_or(&false) {
         return;
     }
 

crates/lsp/src/tower_lsp.rs

@@ -20,6 +20,7 @@ use crate::handlers_ext::ViewFileParams;
 use crate::main_loop::Event;
 use crate::main_loop::GlobalState;
 use crate::main_loop::TokioUnboundedSender;
+use crate::TESTING;
 
 // Based on https://stackoverflow.com/a/69324393/1725177
 macro_rules! cast_response {
@@ -174,13 +175,39 @@ impl LanguageServer for Backend {
     }
 }
 
-pub async fn start_lsp<I, O>(read: I, write: O)
+/// Entry point for the LSP server
+///
+/// Should be called exactly once per process
+pub async fn start_server<I, O>(read: I, write: O)
+where
+    I: AsyncRead + Unpin,
+    O: AsyncWrite,
+{
+    start_server_impl(read, write, false).await
+}
+
+/// Entry point for the test LSP server
+///
+/// Should be called exactly once per process
+pub async fn start_test_server<I, O>(read: I, write: O)
+where
+    I: AsyncRead + Unpin,
+    O: AsyncWrite,
+{
+    start_server_impl(read, write, true).await
+}
+
+async fn start_server_impl<I, O>(read: I, write: O, testing: bool)
 where
     I: AsyncRead + Unpin,
     O: AsyncWrite,
 {
     log::trace!("Starting LSP");
 
+    TESTING
+        .set(testing)
+        .expect("`TESTING` can only be set once.");
+
     let (service, socket) = new_lsp();
     let server = tower_lsp::Server::new(read, write, socket);
     server.serve(service).await;
@@ -214,62 +241,3 @@ fn new_jsonrpc_error(message: String) -> jsonrpc::Error {
         data: None,
     }
 }
-
-#[cfg(test)]
-pub(crate) async fn start_test_client() -> lsp_test::lsp_client::TestClient {
-    lsp_test::lsp_client::TestClient::new(|server_rx, client_tx| async {
-        start_lsp(server_rx, client_tx).await
-    })
-}
-
-#[cfg(test)]
-pub(crate) async fn init_test_client() -> lsp_test::lsp_client::TestClient {
-    let mut client = start_test_client().await;
-
-    client.initialize().await;
-    client.recv_response().await;
-
-    client
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use assert_matches::assert_matches;
-    use tower_lsp::lsp_types;
-
-    #[tests_macros::lsp_test]
-    async fn test_init() {
-        let mut client = start_test_client().await;
-
-        client.initialize().await;
-
-        let value = client.recv_response().await;
-        let value: lsp_types::InitializeResult =
-            serde_json::from_value(value.result().unwrap().clone()).unwrap();
-
-        assert_matches!(
-            value,
-            lsp_types::InitializeResult {
-                capabilities,
-                server_info
-            } => {
-                assert_matches!(capabilities, ServerCapabilities {
-                    position_encoding,
-                    text_document_sync,
-                    ..
-                } => {
-                    assert_eq!(position_encoding, None);
-                    assert_eq!(text_document_sync, Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::INCREMENTAL)));
-                });
-
-                assert_matches!(server_info, Some(ServerInfo { name, version }) => {
-                    assert!(name.contains("Air Language Server"));
-                    assert!(version.is_some());
-                });
-            }
-        );
-
-        client
-    }
-}

crates/lsp/tests/fixtures/mod.rs

@@ -0,0 +1,5 @@
+// For some reason this setup gives false alarms about dead code
+#[allow(dead_code)]
+pub mod test_client;
+#[allow(dead_code)]
+pub mod tower_lsp_test_client;

crates/lsp/tests/fixtures/test_client.rs

@@ -0,0 +1,15 @@
+use lsp::tower_lsp::start_test_server;
+use lsp_test::lsp_client::TestClient;
+
+pub async fn start_test_client() -> lsp_test::lsp_client::TestClient {
+    TestClient::new(|server_rx, client_tx| async { start_test_server(server_rx, client_tx).await })
+}
+
+pub async fn init_test_client() -> TestClient {
+    let mut client = start_test_client().await;
+
+    client.initialize().await;
+    client.recv_response().await;
+
+    client
+}

crates/lsp/src/tower_lsp_test_client.rs → crates/lsp/tests/fixtures/tower_lsp_test_client.rs

@@ -1,7 +1,8 @@
 use lsp_test::lsp_client::TestClient;
 use tower_lsp::lsp_types;
 
-use crate::{documents::Document, from_proto};
+use lsp::documents::Document;
+use lsp::from_proto;
 
 pub(crate) trait TestClientExt {
     async fn open_document(&mut self, doc: &Document) -> lsp_types::TextDocumentItem;

crates/lsp/tests/format_diff.rs

@@ -0,0 +1,27 @@
+mod fixtures;
+
+use fixtures::test_client::init_test_client;
+use fixtures::tower_lsp_test_client::TestClientExt;
+use lsp::documents::Document;
+
+// https://github.com/posit-dev/air/issues/61
+#[tests_macros::lsp_test]
+async fn test_format_minimal_diff() {
+    let mut client = init_test_client().await;
+
+    #[rustfmt::skip]
+    let doc = Document::doodle(
+"1
+2+2
+3
+",
+    );
+
+    let edits = client.format_document_edits(&doc).await.unwrap();
+    assert!(edits.len() == 1);
+
+    let edit = &edits[0];
+    assert_eq!(edit.new_text, " + ");
+
+    client
+}

crates/lsp/tests/format_document.rs

@@ -0,0 +1,24 @@
+mod fixtures;
+
+use fixtures::test_client::init_test_client;
+use fixtures::tower_lsp_test_client::TestClientExt;
+use lsp::documents::Document;
+
+#[tests_macros::lsp_test]
+async fn test_format() {
+    let mut client = init_test_client().await;
+
+    #[rustfmt::skip]
+    let doc = Document::doodle(
+"
+1
+2+2
+3 + 3 +
+3",
+    );
+
+    let formatted = client.format_document(&doc).await;
+    insta::assert_snapshot!(formatted);
+
+    client
+}