[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

@@ -13,6 +13,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 +33,7 @@ pub fn run_server() -> anyhow::Result<()> {
 
     let (connection, io_threads) = Connection::stdio();
 
-    let server_result = Server::new(worker_threads, connection)
+    let server_result = Server::new(worker_threads, connection, None)
         .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,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: Option<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,7 +107,11 @@ impl Server {
                     .collect()
             })
             .or_else(|| {
-                let current_dir = std::env::current_dir().ok()?;
+                let current_dir = if let Some(ref system) = fallback_system {
+                    system.current_directory().as_std_path().to_path_buf()
+                } else {
+                    std::env::current_dir().ok()?
+                };
                 tracing::warn!(
                     "No workspace(s) were provided during initialization. \
                     Using the current working directory as a default workspace: {}",
@@ -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: Option<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: Option<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.as_ref())
+                })
+            }
             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.as_ref())
+                }),
             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: Option<&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.cloned());
             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: Option<&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.

crates/ty_server/src/session.rs

@@ -48,7 +48,7 @@ struct WorkspaceOptions {
 
 /// This is a direct representation of the settings schema sent by the client.
 #[derive(Clone, Debug, Deserialize, Default)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
+#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct ClientOptions {
     /// Settings under the `python.*` namespace in VS Code that are useful for the ty language
@@ -62,7 +62,7 @@ pub(crate) struct ClientOptions {
 
 /// Diagnostic mode for the language server.
 #[derive(Clone, Copy, Debug, Default, Deserialize)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
+#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
 #[serde(rename_all = "camelCase")]
 pub(crate) enum DiagnosticMode {
     /// Check only currently open files.
@@ -139,21 +139,21 @@ impl ClientOptions {
 // all settings and not just the ones in "python.*".
 
 #[derive(Clone, Debug, Deserialize, Default)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
+#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
 #[serde(rename_all = "camelCase")]
 struct Python {
     ty: Option<Ty>,
 }
 
 #[derive(Clone, Debug, Deserialize, Default)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
+#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
 #[serde(rename_all = "camelCase")]
 struct PythonExtension {
     active_environment: Option<ActiveEnvironment>,
 }
 
 #[derive(Clone, Debug, Deserialize)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
+#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct ActiveEnvironment {
     pub(crate) executable: PythonExecutable,
@@ -162,7 +162,7 @@ pub(crate) struct ActiveEnvironment {
 }
 
 #[derive(Clone, Debug, Deserialize)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
+#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct EnvironmentVersion {
     pub(crate) major: i64,
@@ -174,7 +174,7 @@ pub(crate) struct EnvironmentVersion {
 }
 
 #[derive(Clone, Debug, Deserialize)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
+#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct PythonEnvironment {
     pub(crate) folder_uri: Url,
@@ -186,7 +186,7 @@ pub(crate) struct PythonEnvironment {
 }
 
 #[derive(Clone, Debug, Deserialize)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
+#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct PythonExecutable {
     #[allow(dead_code)]
@@ -195,7 +195,7 @@ pub(crate) struct PythonExecutable {
 }
 
 #[derive(Clone, Debug, Deserialize, Default)]
-#[cfg_attr(test, derive(PartialEq, Eq))]
+#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
 #[serde(rename_all = "camelCase")]
 struct Ty {
     disable_language_services: Option<bool>,