[ty] Implement mock language server for testing (#19391)

## Summary

Closes: astral-sh/ty#88

This PR implements an initial version of a mock language server that can
be used to write e2e tests using the real server running in the
background.

The way it works is that you'd use the `TestServerBuilder` to help
construct the `TestServer` with the setup data. This could be the
workspace folders, populating the file and it's content in the memory
file system, setting the right client capabilities to make the server
respond correctly, etc. This can be expanded as we write more test
cases.

There are still a few things to follow-up on:
- ~In the `Drop` implementation, we should assert that there are no
pending notification, request and responses from the server that the
test code hasn't handled yet~ Implemented in [`afd1f82`
(#19391)](afd1f82)
- Reduce the setup boilerplate in any way we can
- Improve the final assertion, currently I'm just snapshotting the final
output

## Test Plan

Written a few test cases.

Author: dhruvmanila

Cargo.lock

@@ -38,6 +38,10 @@ tracing = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["chrono"] }
 
 [dev-dependencies]
+dunce = { workspace = true }
+insta = { workspace = true, features = ["filters", "json"] }
+regex = { workspace = true }
+tempfile = { workspace = true }
 
 [target.'cfg(target_vendor = "apple")'.dependencies]
 libc = { workspace = true }

crates/ty_server/Cargo.toml

@@ -1,7 +1,8 @@
-use std::num::NonZeroUsize;
+use std::{num::NonZeroUsize, sync::Arc};
 
 use anyhow::Context;
 use lsp_server::Connection;
+use ruff_db::system::{OsSystem, SystemPathBuf};
 
 use crate::server::Server;
 pub use document::{NotebookDocument, PositionEncoding, TextDocument};
@@ -13,6 +14,9 @@ mod server;
 mod session;
 mod system;
 
+#[cfg(test)]
+pub mod test;
+
 pub(crate) const SERVER_NAME: &str = "ty";
 pub(crate) const DIAGNOSTIC_NAME: &str = "ty";
 
@@ -30,7 +34,21 @@ pub fn run_server() -> anyhow::Result<()> {
 
     let (connection, io_threads) = Connection::stdio();
 
-    let server_result = Server::new(worker_threads, connection)
+    let cwd = {
+        let cwd = std::env::current_dir().context("Failed to get the current working directory")?;
+        SystemPathBuf::from_path_buf(cwd).map_err(|path| {
+            anyhow::anyhow!(
+                "The current working directory `{}` contains non-Unicode characters. \
+                    ty only supports Unicode paths.",
+                path.display()
+            )
+        })?
+    };
+
+    // This is to complement the `LSPSystem` if the document is not available in the index.
+    let fallback_system = Arc::new(OsSystem::new(cwd));
+
+    let server_result = Server::new(worker_threads, connection, fallback_system, true)
         .context("Failed to start server")?
         .run();
 

crates/ty_server/src/lib.rs

@@ -11,8 +11,9 @@ use lsp_types::{
     ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
     TextDocumentSyncOptions, TypeDefinitionProviderCapability, Url, WorkDoneProgressOptions,
 };
+use ruff_db::system::System;
 use std::num::NonZeroUsize;
-use std::panic::PanicHookInfo;
+use std::panic::{PanicHookInfo, RefUnwindSafe};
 use std::sync::Arc;
 
 mod api;
@@ -35,7 +36,12 @@ pub(crate) struct Server {
 }
 
 impl Server {
-    pub(crate) fn new(worker_threads: NonZeroUsize, connection: Connection) -> crate::Result<Self> {
+    pub(crate) fn new(
+        worker_threads: NonZeroUsize,
+        connection: Connection,
+        native_system: Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
+        initialize_logging: bool,
+    ) -> crate::Result<Self> {
         let (id, init_value) = connection.initialize_start()?;
         let init_params: InitializeParams = serde_json::from_value(init_value)?;
 
@@ -71,10 +77,12 @@ impl Server {
         let (main_loop_sender, main_loop_receiver) = crossbeam::channel::bounded(32);
         let client = Client::new(main_loop_sender.clone(), connection.sender.clone());
 
-        crate::logging::init_logging(
-            global_options.tracing.log_level.unwrap_or_default(),
-            global_options.tracing.log_file.as_deref(),
-        );
+        if initialize_logging {
+            crate::logging::init_logging(
+                global_options.tracing.log_level.unwrap_or_default(),
+                global_options.tracing.log_file.as_deref(),
+            );
+        }
 
         tracing::debug!("Version: {version}");
 
@@ -102,10 +110,14 @@ impl Server {
                     .collect()
             })
             .or_else(|| {
-                let current_dir = std::env::current_dir().ok()?;
+                let current_dir = native_system
+                    .current_directory()
+                    .as_std_path()
+                    .to_path_buf();
                 tracing::warn!(
                     "No workspace(s) were provided during initialization. \
-                    Using the current working directory as a default workspace: {}",
+                    Using the current working directory from the fallback system as a \
+                    default workspace: {}",
                     current_dir.display()
                 );
                 let uri = Url::from_file_path(current_dir).ok()?;
@@ -143,6 +155,7 @@ impl Server {
                 position_encoding,
                 global_options,
                 workspaces,
+                native_system,
             )?,
             client_capabilities,
         })
@@ -288,3 +301,89 @@ impl Drop for ServerPanicHookHandler {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use lsp_types::notification::PublishDiagnostics;
+    use ruff_db::system::SystemPath;
+
+    use crate::session::ClientOptions;
+    use crate::test::TestServerBuilder;
+
+    #[test]
+    fn initialization() -> Result<()> {
+        let server = TestServerBuilder::new()?
+            .build()?
+            .wait_until_workspaces_are_initialized()?;
+
+        let initialization_result = server.initialization_result().unwrap();
+
+        insta::assert_json_snapshot!("initialization", initialization_result);
+
+        Ok(())
+    }
+
+    #[test]
+    fn initialization_with_workspace() -> Result<()> {
+        let workspace_root = SystemPath::new("foo");
+        let server = TestServerBuilder::new()?
+            .with_workspace(workspace_root, ClientOptions::default())?
+            .build()?
+            .wait_until_workspaces_are_initialized()?;
+
+        let initialization_result = server.initialization_result().unwrap();
+
+        insta::assert_json_snapshot!("initialization_with_workspace", initialization_result);
+
+        Ok(())
+    }
+
+    #[test]
+    fn publish_diagnostics_on_did_open() -> Result<()> {
+        let workspace_root = SystemPath::new("src");
+        let foo = SystemPath::new("src/foo.py");
+        let foo_content = "\
+def foo() -> str:
+    return 42
+";
+
+        let mut server = TestServerBuilder::new()?
+            .with_workspace(workspace_root, ClientOptions::default())?
+            .with_file(foo, foo_content)?
+            .enable_pull_diagnostics(false)
+            .build()?
+            .wait_until_workspaces_are_initialized()?;
+
+        server.open_text_document(foo, &foo_content, 1);
+        let diagnostics = server.await_notification::<PublishDiagnostics>()?;
+
+        insta::assert_debug_snapshot!(diagnostics);
+
+        Ok(())
+    }
+
+    #[test]
+    fn pull_diagnostics_on_did_open() -> Result<()> {
+        let workspace_root = SystemPath::new("src");
+        let foo = SystemPath::new("src/foo.py");
+        let foo_content = "\
+def foo() -> str:
+    return 42
+";
+
+        let mut server = TestServerBuilder::new()?
+            .with_workspace(workspace_root, ClientOptions::default())?
+            .with_file(foo, foo_content)?
+            .enable_pull_diagnostics(true)
+            .build()?
+            .wait_until_workspaces_are_initialized()?;
+
+        server.open_text_document(foo, &foo_content, 1);
+        let diagnostics = server.document_diagnostic_request(foo)?;
+
+        insta::assert_debug_snapshot!(diagnostics);
+
+        Ok(())
+    }
+}

crates/ty_server/src/server.rs

@@ -2,11 +2,14 @@
 
 use std::collections::{BTreeMap, VecDeque};
 use std::ops::{Deref, DerefMut};
+use std::panic::RefUnwindSafe;
 use std::sync::Arc;
 
 use anyhow::{Context, anyhow};
 use index::DocumentQueryError;
 use lsp_server::Message;
+use lsp_types::notification::{Exit, Notification};
+use lsp_types::request::{Request, Shutdown};
 use lsp_types::{ClientCapabilities, TextDocumentContentChangeEvent, Url};
 use options::GlobalOptions;
 use ruff_db::Db;
@@ -37,6 +40,9 @@ mod settings;
 
 /// The global state for the LSP
 pub(crate) struct Session {
+    /// A native system to use with the [`LSPSystem`].
+    native_system: Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
+
     /// Used to retrieve information about open documents and settings.
     ///
     /// This will be [`None`] when a mutable reference is held to the index via [`index_mut`]
@@ -99,6 +105,7 @@ impl Session {
         position_encoding: PositionEncoding,
         global_options: GlobalOptions,
         workspace_folders: Vec<(Url, ClientOptions)>,
+        native_system: Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
     ) -> crate::Result<Self> {
         let index = Arc::new(Index::new(global_options.into_settings()));
 
@@ -108,6 +115,7 @@ impl Session {
         }
 
         Ok(Self {
+            native_system,
             position_encoding,
             workspaces,
             deferred_messages: VecDeque::new(),
@@ -155,6 +163,9 @@ impl Session {
         } else {
             match &message {
                 Message::Request(request) => {
+                    if request.method == Shutdown::METHOD {
+                        return Some(message);
+                    }
                     tracing::debug!(
                         "Deferring `{}` request until all workspaces are initialized",
                         request.method
@@ -165,6 +176,9 @@ impl Session {
                     return Some(message);
                 }
                 Message::Notification(notification) => {
+                    if notification.method == Exit::METHOD {
+                        return Some(message);
+                    }
                     tracing::debug!(
                         "Deferring `{}` notification until all workspaces are initialized",
                         notification.method
@@ -218,9 +232,12 @@ impl Session {
     /// If the path is a virtual path, it will return the first project database in the session.
     pub(crate) fn project_state(&self, path: &AnySystemPath) -> &ProjectState {
         match path {
-            AnySystemPath::System(system_path) => self
-                .project_state_for_path(system_path)
-                .unwrap_or_else(|| self.default_project.get(self.index.as_ref())),
+            AnySystemPath::System(system_path) => {
+                self.project_state_for_path(system_path).unwrap_or_else(|| {
+                    self.default_project
+                        .get(self.index.as_ref(), &self.native_system)
+                })
+            }
             AnySystemPath::SystemVirtual(_virtual_path) => {
                 // TODO: Currently, ty only supports single workspace but we need to figure out
                 // which project should this virtual path belong to when there are multiple
@@ -247,7 +264,10 @@ impl Session {
                 .range_mut(..=system_path.to_path_buf())
                 .next_back()
                 .map(|(_, project)| project)
-                .unwrap_or_else(|| self.default_project.get_mut(self.index.as_ref())),
+                .unwrap_or_else(|| {
+                    self.default_project
+                        .get_mut(self.index.as_ref(), &self.native_system)
+                }),
             AnySystemPath::SystemVirtual(_virtual_path) => {
                 // TODO: Currently, ty only supports single workspace but we need to figure out
                 // which project should this virtual path belong to when there are multiple
@@ -330,7 +350,10 @@ impl Session {
             // For now, create one project database per workspace.
             // In the future, index the workspace directories to find all projects
             // and create a project database for each.
-            let system = LSPSystem::new(self.index.as_ref().unwrap().clone());
+            let system = LSPSystem::new(
+                self.index.as_ref().unwrap().clone(),
+                self.native_system.clone(),
+            );
 
             let project = ProjectMetadata::discover(&root, &system)
                 .context("Failed to discover project configuration")
@@ -748,12 +771,16 @@ impl DefaultProject {
         DefaultProject(std::sync::OnceLock::new())
     }
 
-    pub(crate) fn get(&self, index: Option<&Arc<Index>>) -> &ProjectState {
+    pub(crate) fn get(
+        &self,
+        index: Option<&Arc<Index>>,
+        fallback_system: &Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
+    ) -> &ProjectState {
         self.0.get_or_init(|| {
             tracing::info!("Initializing the default project");
 
             let index = index.unwrap();
-            let system = LSPSystem::new(index.clone());
+            let system = LSPSystem::new(index.clone(), fallback_system.clone());
             let metadata = ProjectMetadata::from_options(
                 Options::default(),
                 system.current_directory().to_path_buf(),
@@ -771,8 +798,12 @@ impl DefaultProject {
         })
     }
 
-    pub(crate) fn get_mut(&mut self, index: Option<&Arc<Index>>) -> &mut ProjectState {
-        let _ = self.get(index);
+    pub(crate) fn get_mut(
+        &mut self,
+        index: Option<&Arc<Index>>,
+        fallback_system: &Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
+    ) -> &mut ProjectState {
+        let _ = self.get(index, fallback_system);
 
         // SAFETY: The `OnceLock` is guaranteed to be initialized at this point because
         // we called `get` above, which initializes it if it wasn't already.