[ty] Implement mock language server for testing

crates/ty_server/Cargo.toml

@@ -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/src/lib.rs

@@ -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/server.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/session.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.