[ty] Implement mock language server for testing

Author: dhruvmanila

Cargo.lock

@@ -37,6 +37,7 @@ tracing = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["chrono"] }
 
 [dev-dependencies]
+insta = { workspace = true, features = ["json"] }
 
 [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)
         .context("Failed to start server")?
         .run();
 

crates/ty_server/src/lib.rs

@@ -4,7 +4,7 @@
 //! are written to `stderr` by default, which should appear in the logs for most LSP clients. A
 //! `logFile` path can also be specified in the settings, and output will be directed there
 //! instead.
-use std::sync::Arc;
+use std::sync::{Arc, Once};
 
 use ruff_db::system::{SystemPath, SystemPathBuf};
 use serde::Deserialize;
@@ -14,51 +14,55 @@ use tracing_subscriber::fmt::time::ChronoLocal;
 use tracing_subscriber::fmt::writer::BoxMakeWriter;
 use tracing_subscriber::layer::SubscriberExt;
 
+static INIT_LOGGING: Once = Once::new();
+
 pub(crate) fn init_logging(log_level: LogLevel, log_file: Option<&SystemPath>) {
-    let log_file = log_file
-        .map(|path| {
-            // this expands `logFile` so that tildes and environment variables
-            // are replaced with their values, if possible.
-            if let Some(expanded) = shellexpand::full(&path.to_string())
-                .ok()
-                .map(|path| SystemPathBuf::from(&*path))
-            {
-                expanded
-            } else {
-                path.to_path_buf()
-            }
-        })
-        .and_then(|path| {
-            std::fs::OpenOptions::new()
-                .create(true)
-                .append(true)
-                .open(path.as_std_path())
-                .map_err(|err| {
-                    #[expect(clippy::print_stderr)]
-                    {
-                        eprintln!("Failed to open file at {path} for logging: {err}");
-                    }
-                })
-                .ok()
-        });
+    INIT_LOGGING.call_once(|| {
+        let log_file = log_file
+            .map(|path| {
+                // this expands `logFile` so that tildes and environment variables
+                // are replaced with their values, if possible.
+                if let Some(expanded) = shellexpand::full(&path.to_string())
+                    .ok()
+                    .map(|path| SystemPathBuf::from(&*path))
+                {
+                    expanded
+                } else {
+                    path.to_path_buf()
+                }
+            })
+            .and_then(|path| {
+                std::fs::OpenOptions::new()
+                    .create(true)
+                    .append(true)
+                    .open(path.as_std_path())
+                    .map_err(|err| {
+                        #[expect(clippy::print_stderr)]
+                        {
+                            eprintln!("Failed to open file at {path} for logging: {err}");
+                        }
+                    })
+                    .ok()
+            });
 
-    let logger = match log_file {
-        Some(file) => BoxMakeWriter::new(Arc::new(file)),
-        None => BoxMakeWriter::new(std::io::stderr),
-    };
-    let is_trace_level = log_level == LogLevel::Trace;
-    let subscriber = tracing_subscriber::Registry::default().with(
-        tracing_subscriber::fmt::layer()
-            .with_timer(ChronoLocal::new("%Y-%m-%d %H:%M:%S.%f".to_string()))
-            .with_thread_names(is_trace_level)
-            .with_target(is_trace_level)
-            .with_ansi(false)
-            .with_writer(logger)
-            .with_filter(LogLevelFilter { filter: log_level }),
-    );
+        let logger = match log_file {
+            Some(file) => BoxMakeWriter::new(Arc::new(file)),
+            None => BoxMakeWriter::new(std::io::stderr),
+        };
+        let is_trace_level = log_level == LogLevel::Trace;
+        let subscriber = tracing_subscriber::Registry::default().with(
+            tracing_subscriber::fmt::layer()
+                .with_timer(ChronoLocal::new("%Y-%m-%d %H:%M:%S.%f".to_string()))
+                .with_thread_names(is_trace_level)
+                .with_target(is_trace_level)
+                .with_ansi(false)
+                .with_writer(logger)
+                .with_filter(LogLevelFilter { filter: log_level }),
+        );
 
-    tracing::subscriber::set_global_default(subscriber)
-        .expect("should be able to set global default subscriber");
+        tracing::subscriber::set_global_default(subscriber)
+            .expect("should be able to set global default subscriber");
+    });
 }
 
 /// The log level for the server as provided by the client during initialization.

crates/ty_server/src/logging.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,11 @@ 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,
+        fallback_system: Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
+    ) -> crate::Result<Self> {
         let (id, init_value) = connection.initialize_start()?;
         let init_params: InitializeParams = serde_json::from_value(init_value)?;
 
@@ -102,10 +107,14 @@ impl Server {
                     .collect()
             })
             .or_else(|| {
-                let current_dir = std::env::current_dir().ok()?;
+                let current_dir = fallback_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 +152,7 @@ impl Server {
                 position_encoding,
                 global_options,
                 workspaces,
+                fallback_system,
             )?,
             client_capabilities,
         })
@@ -286,3 +296,43 @@ impl Drop for ServerPanicHookHandler {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use ruff_db::system::{InMemorySystem, SystemPathBuf};
+
+    use crate::session::ClientOptions;
+    use crate::test::TestServerBuilder;
+
+    #[test]
+    fn initialization_sequence() {
+        let system = InMemorySystem::default();
+        let test_server = TestServerBuilder::new()
+            .with_memory_system(system)
+            .build()
+            .unwrap()
+            .wait_until_workspaces_are_initialized()
+            .unwrap();
+
+        let initialization_result = test_server.initialization_result().unwrap();
+
+        insta::assert_json_snapshot!("initialization_capabilities", initialization_result);
+    }
+
+    #[test]
+    fn initialization_with_workspace() {
+        let workspace_root = SystemPathBuf::from("/foo");
+        let system = InMemorySystem::new(workspace_root.clone());
+        let test_server = TestServerBuilder::new()
+            .with_memory_system(system)
+            .with_workspace(&workspace_root, ClientOptions::default())
+            .build()
+            .unwrap()
+            .wait_until_workspaces_are_initialized()
+            .unwrap();
+
+        let initialization_result = test_server.initialization_result().unwrap();
+
+        insta::assert_json_snapshot!("initialization_with_workspace", initialization_result);
+    }
+}

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;
@@ -36,6 +39,9 @@ mod settings;
 
 /// The global state for the LSP
 pub(crate) struct Session {
+    /// A fallback system to use with the [`LSPSystem`].
+    fallback_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`]
@@ -79,6 +85,7 @@ impl Session {
         position_encoding: PositionEncoding,
         global_options: GlobalOptions,
         workspace_folders: Vec<(Url, ClientOptions)>,
+        fallback_system: Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
     ) -> crate::Result<Self> {
         let index = Arc::new(Index::new(global_options.into_settings()));
 
@@ -88,6 +95,7 @@ impl Session {
         }
 
         Ok(Self {
+            fallback_system,
             position_encoding,
             workspaces,
             deferred_messages: VecDeque::new(),
@@ -137,6 +145,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
@@ -147,6 +158,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
@@ -171,9 +185,12 @@ impl Session {
     /// If the path is a virtual path, it will return the first project database in the session.
     pub(crate) fn project_db(&self, path: &AnySystemPath) -> &ProjectDatabase {
         match path {
-            AnySystemPath::System(system_path) => self
-                .project_db_for_path(system_path)
-                .unwrap_or_else(|| self.default_project.get(self.index.as_ref())),
+            AnySystemPath::System(system_path) => {
+                self.project_db_for_path(system_path).unwrap_or_else(|| {
+                    self.default_project
+                        .get(self.index.as_ref(), &self.fallback_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
@@ -196,7 +213,10 @@ impl Session {
                 .range_mut(..=system_path.to_path_buf())
                 .next_back()
                 .map(|(_, db)| db)
-                .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.fallback_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
@@ -268,7 +288,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.fallback_system.clone(),
+            );
 
             let project = ProjectMetadata::discover(&root, &system)
                 .context("Failed to discover project configuration")
@@ -663,11 +686,15 @@ impl DefaultProject {
         DefaultProject(std::sync::OnceLock::new())
     }
 
-    pub(crate) fn get(&self, index: Option<&Arc<Index>>) -> &ProjectDatabase {
+    pub(crate) fn get(
+        &self,
+        index: Option<&Arc<Index>>,
+        fallback_system: &Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
+    ) -> &ProjectDatabase {
         self.0.get_or_init(|| {
             tracing::info!("Initialize default project");
 
-            let system = LSPSystem::new(index.unwrap().clone());
+            let system = LSPSystem::new(index.unwrap().clone(), fallback_system.clone());
             let metadata = ProjectMetadata::from_options(
                 Options::default(),
                 system.current_directory().to_path_buf(),
@@ -678,8 +705,12 @@ impl DefaultProject {
         })
     }
 
-    pub(crate) fn get_mut(&mut self, index: Option<&Arc<Index>>) -> &mut ProjectDatabase {
-        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 ProjectDatabase {
+        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.