[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)](afd1f82bde)
- 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.

Cargo.lock

@@ -4322,10 +4322,13 @@ dependencies = [
  "anyhow",
  "bitflags 2.9.1",
  "crossbeam",
+ "dunce",
+ "insta",
  "jod-thread",
  "libc",
  "lsp-server",
  "lsp-types",
+ "regex",
  "ruff_db",
  "ruff_notebook",
  "ruff_python_ast",
@@ -4336,6 +4339,7 @@ dependencies = [
  "serde",
  "serde_json",
  "shellexpand",
+ "tempfile",
  "thiserror 2.0.12",
  "tracing",
  "tracing-subscriber",

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());
 
+        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.

crates/ty_server/src/session/options.rs

@@ -49,7 +49,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
@@ -63,7 +63,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.
@@ -147,21 +147,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,
@@ -170,7 +170,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,
@@ -182,7 +182,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,
@@ -194,7 +194,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)]
@@ -203,7 +203,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>,

crates/ty_server/src/snapshots/ty_server__server__tests__initialization.snap

@@ -0,0 +1,69 @@
+---
+source: crates/ty_server/src/server.rs
+expression: initialization_result
+---
+{
+  "capabilities": {
+    "positionEncoding": "utf-16",
+    "textDocumentSync": {
+      "openClose": true,
+      "change": 2
+    },
+    "hoverProvider": true,
+    "completionProvider": {
+      "triggerCharacters": [
+        "."
+      ]
+    },
+    "signatureHelpProvider": {
+      "triggerCharacters": [
+        "(",
+        ","
+      ],
+      "retriggerCharacters": [
+        ")"
+      ]
+    },
+    "definitionProvider": true,
+    "typeDefinitionProvider": true,
+    "declarationProvider": true,
+    "semanticTokensProvider": {
+      "legend": {
+        "tokenTypes": [
+          "namespace",
+          "class",
+          "parameter",
+          "selfParameter",
+          "clsParameter",
+          "variable",
+          "property",
+          "function",
+          "method",
+          "keyword",
+          "string",
+          "number",
+          "decorator",
+          "builtinConstant",
+          "typeParameter"
+        ],
+        "tokenModifiers": [
+          "definition",
+          "readonly",
+          "async"
+        ]
+      },
+      "range": true,
+      "full": true
+    },
+    "inlayHintProvider": {},
+    "diagnosticProvider": {
+      "identifier": "ty",
+      "interFileDependencies": true,
+      "workspaceDiagnostics": false
+    }
+  },
+  "serverInfo": {
+    "name": "ty",
+    "version": "Unknown"
+  }
+}

crates/ty_server/src/snapshots/ty_server__server__tests__initialization_with_workspace.snap

@@ -0,0 +1,69 @@
+---
+source: crates/ty_server/src/server.rs
+expression: initialization_result
+---
+{
+  "capabilities": {
+    "positionEncoding": "utf-16",
+    "textDocumentSync": {
+      "openClose": true,
+      "change": 2
+    },
+    "hoverProvider": true,
+    "completionProvider": {
+      "triggerCharacters": [
+        "."
+      ]
+    },
+    "signatureHelpProvider": {
+      "triggerCharacters": [
+        "(",
+        ","
+      ],
+      "retriggerCharacters": [
+        ")"
+      ]
+    },
+    "definitionProvider": true,
+    "typeDefinitionProvider": true,
+    "declarationProvider": true,
+    "semanticTokensProvider": {
+      "legend": {
+        "tokenTypes": [
+          "namespace",
+          "class",
+          "parameter",
+          "selfParameter",
+          "clsParameter",
+          "variable",
+          "property",
+          "function",
+          "method",
+          "keyword",
+          "string",
+          "number",
+          "decorator",
+          "builtinConstant",
+          "typeParameter"
+        ],
+        "tokenModifiers": [
+          "definition",
+          "readonly",
+          "async"
+        ]
+      },
+      "range": true,
+      "full": true
+    },
+    "inlayHintProvider": {},
+    "diagnosticProvider": {
+      "identifier": "ty",
+      "interFileDependencies": true,
+      "workspaceDiagnostics": false
+    }
+  },
+  "serverInfo": {
+    "name": "ty",
+    "version": "Unknown"
+  }
+}

crates/ty_server/src/snapshots/ty_server__server__tests__publish_diagnostics_on_did_open.snap

@@ -0,0 +1,99 @@
+---
+source: crates/ty_server/src/server.rs
+expression: diagnostics
+---
+PublishDiagnosticsParams {
+    uri: Url {
+        scheme: "file",
+        cannot_be_a_base: false,
+        username: "",
+        password: None,
+        host: None,
+        port: None,
+        path: "<temp_dir>/src/foo.py",
+        query: None,
+        fragment: None,
+    },
+    diagnostics: [
+        Diagnostic {
+            range: Range {
+                start: Position {
+                    line: 1,
+                    character: 11,
+                },
+                end: Position {
+                    line: 1,
+                    character: 13,
+                },
+            },
+            severity: Some(
+                Error,
+            ),
+            code: Some(
+                String(
+                    "invalid-return-type",
+                ),
+            ),
+            code_description: Some(
+                CodeDescription {
+                    href: Url {
+                        scheme: "https",
+                        cannot_be_a_base: false,
+                        username: "",
+                        password: None,
+                        host: Some(
+                            Domain(
+                                "ty.dev",
+                            ),
+                        ),
+                        port: None,
+                        path: "/rules",
+                        query: None,
+                        fragment: Some(
+                            "invalid-return-type",
+                        ),
+                    },
+                },
+            ),
+            source: Some(
+                "ty",
+            ),
+            message: "Return type does not match returned value: expected `str`, found `Literal[42]`",
+            related_information: Some(
+                [
+                    DiagnosticRelatedInformation {
+                        location: Location {
+                            uri: Url {
+                                scheme: "file",
+                                cannot_be_a_base: false,
+                                username: "",
+                                password: None,
+                                host: None,
+                                port: None,
+                                path: "<temp_dir>/src/foo.py",
+                                query: None,
+                                fragment: None,
+                            },
+                            range: Range {
+                                start: Position {
+                                    line: 0,
+                                    character: 13,
+                                },
+                                end: Position {
+                                    line: 0,
+                                    character: 16,
+                                },
+                            },
+                        },
+                        message: "Expected `str` because of return type",
+                    },
+                ],
+            ),
+            tags: None,
+            data: None,
+        },
+    ],
+    version: Some(
+        1,
+    ),
+}

crates/ty_server/src/snapshots/ty_server__server__tests__pull_diagnostics_on_did_open.snap

@@ -0,0 +1,93 @@
+---
+source: crates/ty_server/src/server.rs
+expression: diagnostics
+---
+Report(
+    Full(
+        RelatedFullDocumentDiagnosticReport {
+            related_documents: None,
+            full_document_diagnostic_report: FullDocumentDiagnosticReport {
+                result_id: None,
+                items: [
+                    Diagnostic {
+                        range: Range {
+                            start: Position {
+                                line: 1,
+                                character: 11,
+                            },
+                            end: Position {
+                                line: 1,
+                                character: 13,
+                            },
+                        },
+                        severity: Some(
+                            Error,
+                        ),
+                        code: Some(
+                            String(
+                                "invalid-return-type",
+                            ),
+                        ),
+                        code_description: Some(
+                            CodeDescription {
+                                href: Url {
+                                    scheme: "https",
+                                    cannot_be_a_base: false,
+                                    username: "",
+                                    password: None,
+                                    host: Some(
+                                        Domain(
+                                            "ty.dev",
+                                        ),
+                                    ),
+                                    port: None,
+                                    path: "/rules",
+                                    query: None,
+                                    fragment: Some(
+                                        "invalid-return-type",
+                                    ),
+                                },
+                            },
+                        ),
+                        source: Some(
+                            "ty",
+                        ),
+                        message: "Return type does not match returned value: expected `str`, found `Literal[42]`",
+                        related_information: Some(
+                            [
+                                DiagnosticRelatedInformation {
+                                    location: Location {
+                                        uri: Url {
+                                            scheme: "file",
+                                            cannot_be_a_base: false,
+                                            username: "",
+                                            password: None,
+                                            host: None,
+                                            port: None,
+                                            path: "<temp_dir>/src/foo.py",
+                                            query: None,
+                                            fragment: None,
+                                        },
+                                        range: Range {
+                                            start: Position {
+                                                line: 0,
+                                                character: 13,
+                                            },
+                                            end: Position {
+                                                line: 0,
+                                                character: 16,
+                                            },
+                                        },
+                                    },
+                                    message: "Expected `str` because of return type",
+                                },
+                            ],
+                        ),
+                        tags: None,
+                        data: None,
+                    },
+                ],
+            },
+        },
+    ),
+)

crates/ty_server/src/system.rs

@@ -1,6 +1,7 @@
 use std::any::Any;
 use std::fmt;
 use std::fmt::Display;
+use std::panic::RefUnwindSafe;
 use std::sync::Arc;
 
 use lsp_types::Url;
@@ -8,8 +9,8 @@ use ruff_db::file_revision::FileRevision;
 use ruff_db::files::{File, FilePath};
 use ruff_db::system::walk_directory::WalkDirectoryBuilder;
 use ruff_db::system::{
-    CaseSensitivity, DirectoryEntry, FileType, GlobError, Metadata, OsSystem, PatternError, Result,
-    System, SystemPath, SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf, WritableSystem,
+    CaseSensitivity, DirectoryEntry, FileType, GlobError, Metadata, PatternError, Result, System,
+    SystemPath, SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf, WritableSystem,
 };
 use ruff_notebook::{Notebook, NotebookError};
 use ty_python_semantic::Db;
@@ -118,18 +119,21 @@ pub(crate) struct LSPSystem {
     /// [`index_mut`]: crate::Session::index_mut
     index: Option<Arc<Index>>,
 
-    /// A system implementation that uses the local file system.
-    os_system: OsSystem,
+    /// A native system implementation.
+    ///
+    /// This is used to delegate method calls that are not handled by the LSP system. It is also
+    /// used as a fallback when the documents are not found in the LSP index.
+    native_system: Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
 }
 
 impl LSPSystem {
-    pub(crate) fn new(index: Arc<Index>) -> Self {
-        let cwd = std::env::current_dir().unwrap();
-        let os_system = OsSystem::new(SystemPathBuf::from_path_buf(cwd).unwrap());
-
+    pub(crate) fn new(
+        index: Arc<Index>,
+        native_system: Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
+    ) -> Self {
         Self {
             index: Some(index),
-            os_system,
+            native_system,
         }
     }
 
@@ -183,16 +187,16 @@ impl System for LSPSystem {
                 FileType::File,
             ))
         } else {
-            self.os_system.path_metadata(path)
+            self.native_system.path_metadata(path)
         }
     }
 
     fn canonicalize_path(&self, path: &SystemPath) -> Result<SystemPathBuf> {
-        self.os_system.canonicalize_path(path)
+        self.native_system.canonicalize_path(path)
     }
 
     fn path_exists_case_sensitive(&self, path: &SystemPath, prefix: &SystemPath) -> bool {
-        self.os_system.path_exists_case_sensitive(path, prefix)
+        self.native_system.path_exists_case_sensitive(path, prefix)
     }
 
     fn read_to_string(&self, path: &SystemPath) -> Result<String> {
@@ -200,7 +204,7 @@ impl System for LSPSystem {
 
         match document {
             Some(DocumentQuery::Text { document, .. }) => Ok(document.contents().to_string()),
-            _ => self.os_system.read_to_string(path),
+            _ => self.native_system.read_to_string(path),
         }
     }
 
@@ -212,7 +216,7 @@ impl System for LSPSystem {
                 Notebook::from_source_code(document.contents())
             }
             Some(DocumentQuery::Notebook { notebook, .. }) => Ok(notebook.make_ruff_notebook()),
-            None => self.os_system.read_to_notebook(path),
+            None => self.native_system.read_to_notebook(path),
         }
     }
 
@@ -243,26 +247,26 @@ impl System for LSPSystem {
     }
 
     fn current_directory(&self) -> &SystemPath {
-        self.os_system.current_directory()
+        self.native_system.current_directory()
     }
 
     fn user_config_directory(&self) -> Option<SystemPathBuf> {
-        self.os_system.user_config_directory()
+        self.native_system.user_config_directory()
     }
 
     fn cache_dir(&self) -> Option<SystemPathBuf> {
-        self.os_system.cache_dir()
+        self.native_system.cache_dir()
     }
 
     fn read_directory<'a>(
         &'a self,
         path: &SystemPath,
     ) -> Result<Box<dyn Iterator<Item = Result<DirectoryEntry>> + 'a>> {
-        self.os_system.read_directory(path)
+        self.native_system.read_directory(path)
     }
 
     fn walk_directory(&self, path: &SystemPath) -> WalkDirectoryBuilder {
-        self.os_system.walk_directory(path)
+        self.native_system.walk_directory(path)
     }
 
     fn glob(
@@ -272,11 +276,11 @@ impl System for LSPSystem {
         Box<dyn Iterator<Item = std::result::Result<SystemPathBuf, GlobError>> + '_>,
         PatternError,
     > {
-        self.os_system.glob(pattern)
+        self.native_system.glob(pattern)
     }
 
     fn as_writable(&self) -> Option<&dyn WritableSystem> {
-        self.os_system.as_writable()
+        self.native_system.as_writable()
     }
 
     fn as_any(&self) -> &dyn Any {
@@ -288,11 +292,11 @@ impl System for LSPSystem {
     }
 
     fn case_sensitivity(&self) -> CaseSensitivity {
-        self.os_system.case_sensitivity()
+        self.native_system.case_sensitivity()
     }
 
     fn env_var(&self, name: &str) -> std::result::Result<String, std::env::VarError> {
-        self.os_system.env_var(name)
+        self.native_system.env_var(name)
     }
 }
 

crates/ty_server/src/test.rs

@@ -0,0 +1,889 @@
+//! Testing server for the ty language server.
+//!
+//! This module provides mock server infrastructure for testing LSP functionality using
+//! temporary directories on the real filesystem.
+//!
+//! The design is inspired by the Starlark LSP test server but adapted for ty server architecture.
+
+use std::collections::hash_map::Entry;
+use std::collections::{HashMap, VecDeque};
+use std::num::NonZeroUsize;
+use std::sync::{Arc, OnceLock};
+use std::thread::JoinHandle;
+use std::time::Duration;
+use std::{fmt, fs};
+
+use anyhow::{Context, Result, anyhow};
+use crossbeam::channel::RecvTimeoutError;
+use insta::internals::SettingsBindDropGuard;
+use lsp_server::{Connection, Message, RequestId, Response, ResponseError};
+use lsp_types::notification::{
+    DidChangeTextDocument, DidChangeWatchedFiles, DidCloseTextDocument, DidOpenTextDocument, Exit,
+    Initialized, Notification,
+};
+use lsp_types::request::{
+    DocumentDiagnosticRequest, Initialize, Request, Shutdown, WorkspaceConfiguration,
+};
+use lsp_types::{
+    ClientCapabilities, ConfigurationParams, DiagnosticClientCapabilities,
+    DidChangeTextDocumentParams, DidChangeWatchedFilesClientCapabilities,
+    DidChangeWatchedFilesParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
+    DocumentDiagnosticParams, DocumentDiagnosticReportResult, FileEvent, InitializeParams,
+    InitializeResult, InitializedParams, PartialResultParams, PublishDiagnosticsClientCapabilities,
+    TextDocumentClientCapabilities, TextDocumentContentChangeEvent, TextDocumentIdentifier,
+    TextDocumentItem, Url, VersionedTextDocumentIdentifier, WorkDoneProgressParams,
+    WorkspaceClientCapabilities, WorkspaceFolder,
+};
+use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf, TestSystem};
+use rustc_hash::FxHashMap;
+use serde::de::DeserializeOwned;
+use tempfile::TempDir;
+
+use crate::logging::{LogLevel, init_logging};
+use crate::server::Server;
+use crate::session::ClientOptions;
+
+/// Number of times to retry receiving a message before giving up
+const RETRY_COUNT: usize = 5;
+
+static INIT_TRACING: OnceLock<()> = OnceLock::new();
+
+/// Setup tracing for the test server.
+///
+/// This will make sure that the tracing subscriber is initialized only once, so that running
+/// multiple tests does not cause multiple subscribers to be registered.
+fn setup_tracing() {
+    INIT_TRACING.get_or_init(|| {
+        init_logging(LogLevel::Debug, None);
+    });
+}
+
+/// Errors that can occur during testing
+#[derive(thiserror::Error, Debug)]
+pub(crate) enum TestServerError {
+    /// The response came back, but was an error response, not a successful one.
+    #[error("Response error: {0:?}")]
+    ResponseError(ResponseError),
+
+    #[error("Invalid response message for request {0}: {1:?}")]
+    InvalidResponse(RequestId, Box<Response>),
+
+    #[error("Got a duplicate response for request ID {0}: {1:?}")]
+    DuplicateResponse(RequestId, Box<Response>),
+
+    #[error("Failed to receive message from server: {0}")]
+    RecvTimeoutError(RecvTimeoutError),
+}
+
+impl TestServerError {
+    fn is_disconnected(&self) -> bool {
+        matches!(
+            self,
+            TestServerError::RecvTimeoutError(RecvTimeoutError::Disconnected)
+        )
+    }
+}
+
+/// A test server for the ty language server that provides helpers for sending requests,
+/// correlating responses, and handling notifications.
+///
+/// The [`Drop`] implementation ensures that the server is shut down gracefully using the described
+/// protocol in the LSP specification. It also ensures that all messages sent by the server have
+/// been handled by the test client before the server is dropped.
+pub(crate) struct TestServer {
+    /// The thread that's actually running the server.
+    ///
+    /// This is an [`Option`] so that the join handle can be taken out when the server is dropped,
+    /// allowing the server thread to be joined and cleaned up properly.
+    server_thread: Option<JoinHandle<()>>,
+
+    /// Connection to communicate with the server.
+    ///
+    /// This is an [`Option`] so that it can be taken out when the server is dropped, allowing
+    /// the connection to be cleaned up properly.
+    client_connection: Option<Connection>,
+
+    /// Test directory that holds all test files.
+    ///
+    /// This directory is automatically cleaned up when the [`TestServer`] is dropped.
+    test_dir: TestContext,
+
+    /// Incrementing counter to automatically generate request IDs
+    request_counter: i32,
+
+    /// A mapping of request IDs to responses received from the server
+    responses: FxHashMap<RequestId, Response>,
+
+    /// An ordered queue of all the notifications received from the server
+    notifications: VecDeque<lsp_server::Notification>,
+
+    /// An ordered queue of all the requests received from the server
+    requests: VecDeque<lsp_server::Request>,
+
+    /// The response from server initialization
+    initialize_response: Option<InitializeResult>,
+
+    /// Workspace configurations for `workspace/configuration` requests
+    workspace_configurations: HashMap<Url, ClientOptions>,
+
+    /// Capabilities registered by the server
+    registered_capabilities: Vec<String>,
+}
+
+impl TestServer {
+    /// Create a new test server with the given workspace configurations
+    fn new(
+        workspaces: Vec<(WorkspaceFolder, ClientOptions)>,
+        test_dir: TestContext,
+        capabilities: ClientCapabilities,
+    ) -> Result<Self> {
+        setup_tracing();
+
+        let (server_connection, client_connection) = Connection::memory();
+
+        // Create OS system with the test directory as cwd
+        let os_system = OsSystem::new(test_dir.root());
+
+        // Start the server in a separate thread
+        let server_thread = std::thread::spawn(move || {
+            // TODO: This should probably be configurable to test concurrency issues
+            let worker_threads = NonZeroUsize::new(1).unwrap();
+            let test_system = Arc::new(TestSystem::new(os_system));
+
+            match Server::new(worker_threads, server_connection, test_system, false) {
+                Ok(server) => {
+                    if let Err(err) = server.run() {
+                        panic!("Server stopped with error: {err:?}");
+                    }
+                }
+                Err(err) => {
+                    panic!("Failed to create server: {err:?}");
+                }
+            }
+        });
+
+        let workspace_folders = workspaces
+            .iter()
+            .map(|(folder, _)| folder.clone())
+            .collect::<Vec<_>>();
+
+        let workspace_configurations = workspaces
+            .into_iter()
+            .map(|(folder, options)| (folder.uri, options))
+            .collect::<HashMap<_, _>>();
+
+        Self {
+            server_thread: Some(server_thread),
+            client_connection: Some(client_connection),
+            test_dir,
+            request_counter: 0,
+            responses: FxHashMap::default(),
+            notifications: VecDeque::new(),
+            requests: VecDeque::new(),
+            initialize_response: None,
+            workspace_configurations,
+            registered_capabilities: Vec::new(),
+        }
+        .initialize(workspace_folders, capabilities)
+    }
+
+    /// Perform LSP initialization handshake
+    fn initialize(
+        mut self,
+        workspace_folders: Vec<WorkspaceFolder>,
+        capabilities: ClientCapabilities,
+    ) -> Result<Self> {
+        let init_params = InitializeParams {
+            capabilities,
+            workspace_folders: Some(workspace_folders),
+            // TODO: This should be configurable by the test server builder. This might not be
+            // required after client settings are implemented in the server.
+            initialization_options: Some(serde_json::Value::Object(serde_json::Map::new())),
+            ..Default::default()
+        };
+
+        let init_request_id = self.send_request::<Initialize>(init_params);
+        self.initialize_response = Some(self.await_response::<InitializeResult>(init_request_id)?);
+        self.send_notification::<Initialized>(InitializedParams {});
+
+        Ok(self)
+    }
+
+    /// Wait until the server has initialized all workspaces.
+    ///
+    /// This will wait until the client receives a `workspace/configuration` request from the
+    /// server, and handles the request.
+    ///
+    /// This should only be called if the server is expected to send this request.
+    pub(crate) fn wait_until_workspaces_are_initialized(mut self) -> Result<Self> {
+        let (request_id, params) = self.await_request::<WorkspaceConfiguration>()?;
+        self.handle_workspace_configuration_request(request_id, &params)?;
+        Ok(self)
+    }
+
+    /// Drain all messages from the server.
+    fn drain_messages(&mut self) {
+        loop {
+            // Don't wait too long to drain the messages, as this is called in the `Drop`
+            // implementation which happens everytime the test ends.
+            match self.receive(Some(Duration::from_millis(10))) {
+                Ok(()) => {}
+                Err(TestServerError::RecvTimeoutError(_)) => {
+                    // Only break if we have no more messages to process.
+                    break;
+                }
+                Err(err) => {
+                    tracing::error!("Error while draining messages: {err:?}");
+                }
+            }
+        }
+    }
+
+    /// Validate that there are no pending messages from the server.
+    ///
+    /// This should be called before the test server is dropped to ensure that all server messages
+    /// have been properly consumed by the test. If there are any pending messages, this will panic
+    /// with detailed information about what was left unconsumed.
+    fn assert_no_pending_messages(&self) {
+        let mut errors = Vec::new();
+
+        if !self.responses.is_empty() {
+            errors.push(format!("Unclaimed responses: {:#?}", self.responses));
+        }
+
+        if !self.notifications.is_empty() {
+            errors.push(format!(
+                "Unclaimed notifications: {:#?}",
+                self.notifications
+            ));
+        }
+
+        if !self.requests.is_empty() {
+            errors.push(format!("Unclaimed requests: {:#?}", self.requests));
+        }
+
+        assert!(
+            errors.is_empty(),
+            "Test server has pending messages that were not consumed by the test:\n{}",
+            errors.join("\n")
+        );
+    }
+
+    /// Generate a new request ID
+    fn next_request_id(&mut self) -> RequestId {
+        self.request_counter += 1;
+        RequestId::from(self.request_counter)
+    }
+
+    /// Send a message to the server.
+    ///
+    /// # Panics
+    ///
+    /// If the server is still running but the client connection got dropped, or if the server
+    /// exited unexpectedly or panicked.
+    #[track_caller]
+    fn send(&mut self, message: Message) {
+        if self
+            .client_connection
+            .as_ref()
+            .unwrap()
+            .sender
+            .send(message)
+            .is_err()
+        {
+            self.panic_on_server_disconnect();
+        }
+    }
+
+    /// Send a request to the server and return the request ID.
+    ///
+    /// The caller can use this ID to later retrieve the response using [`get_response`].
+    ///
+    /// [`get_response`]: TestServer::get_response
+    pub(crate) fn send_request<R>(&mut self, params: R::Params) -> RequestId
+    where
+        R: Request,
+    {
+        let id = self.next_request_id();
+        let request = lsp_server::Request::new(id.clone(), R::METHOD.to_string(), params);
+        self.send(Message::Request(request));
+        id
+    }
+
+    /// Send a notification to the server.
+    pub(crate) fn send_notification<N>(&mut self, params: N::Params)
+    where
+        N: Notification,
+    {
+        let notification = lsp_server::Notification::new(N::METHOD.to_string(), params);
+        self.send(Message::Notification(notification));
+    }
+
+    /// Wait for a server response corresponding to the given request ID.
+    ///
+    /// This should only be called if a request was already sent to the server via [`send_request`]
+    /// which returns the request ID that should be used here.
+    ///
+    /// This method will remove the response from the internal data structure, so it can only be
+    /// called once per request ID.
+    ///
+    /// [`send_request`]: TestServer::send_request
+    pub(crate) fn await_response<T: DeserializeOwned>(&mut self, id: RequestId) -> Result<T> {
+        loop {
+            if let Some(response) = self.responses.remove(&id) {
+                match response {
+                    Response {
+                        error: None,
+                        result: Some(result),
+                        ..
+                    } => {
+                        return Ok(serde_json::from_value::<T>(result)?);
+                    }
+                    Response {
+                        error: Some(err),
+                        result: None,
+                        ..
+                    } => {
+                        return Err(TestServerError::ResponseError(err).into());
+                    }
+                    response => {
+                        return Err(TestServerError::InvalidResponse(id, Box::new(response)).into());
+                    }
+                }
+            }
+
+            self.receive_or_panic()?;
+        }
+    }
+
+    /// Wait for a notification of the specified type from the server and return its parameters.
+    ///
+    /// The caller should ensure that the server is expected to send this notification type. It
+    /// will keep polling the server for this notification up to 10 times before giving up after
+    /// which it will return an error. It will also return an error if the notification is not
+    /// received within `recv_timeout` duration.
+    ///
+    /// This method will remove the notification from the internal data structure, so it should
+    /// only be called if the notification is expected to be sent by the server.
+    pub(crate) fn await_notification<N: Notification>(&mut self) -> Result<N::Params> {
+        for retry_count in 0..RETRY_COUNT {
+            if retry_count > 0 {
+                tracing::info!("Retrying to receive `{}` notification", N::METHOD);
+            }
+            let notification = self
+                .notifications
+                .iter()
+                .position(|notification| N::METHOD == notification.method)
+                .and_then(|index| self.notifications.remove(index));
+            if let Some(notification) = notification {
+                return Ok(serde_json::from_value(notification.params)?);
+            }
+            self.receive_or_panic()?;
+        }
+        Err(anyhow::anyhow!(
+            "Failed to receive `{}` notification after {RETRY_COUNT} retries",
+            N::METHOD
+        ))
+    }
+
+    /// Wait for a request of the specified type from the server and return the request ID and
+    /// parameters.
+    ///
+    /// The caller should ensure that the server is expected to send this request type. It will
+    /// keep polling the server for this request up to 10 times before giving up after which it
+    /// will return an error. It can also return an error if the request is not received within
+    /// `recv_timeout` duration.
+    ///
+    /// This method will remove the request from the internal data structure, so it should only be
+    /// called if the request is expected to be sent by the server.
+    pub(crate) fn await_request<R: Request>(&mut self) -> Result<(RequestId, R::Params)> {
+        for retry_count in 0..RETRY_COUNT {
+            if retry_count > 0 {
+                tracing::info!("Retrying to receive `{}` request", R::METHOD);
+            }
+            let request = self
+                .requests
+                .iter()
+                .position(|request| R::METHOD == request.method)
+                .and_then(|index| self.requests.remove(index));
+            if let Some(request) = request {
+                let params = serde_json::from_value(request.params)?;
+                return Ok((request.id, params));
+            }
+            self.receive_or_panic()?;
+        }
+        Err(anyhow::anyhow!(
+            "Failed to receive `{}` request after {RETRY_COUNT} retries",
+            R::METHOD
+        ))
+    }
+
+    /// Receive a message from the server.
+    ///
+    /// It will wait for `timeout` duration for a message to arrive. If no message is received
+    /// within that time, it will return an error.
+    ///
+    /// If `timeout` is `None`, it will use a default timeout of 1 second.
+    fn receive(&mut self, timeout: Option<Duration>) -> Result<(), TestServerError> {
+        static DEFAULT_TIMEOUT: Duration = Duration::from_secs(1);
+
+        let receiver = self.client_connection.as_ref().unwrap().receiver.clone();
+        let message = receiver
+            .recv_timeout(timeout.unwrap_or(DEFAULT_TIMEOUT))
+            .map_err(TestServerError::RecvTimeoutError)?;
+
+        self.handle_message(message)?;
+
+        for message in receiver.try_iter() {
+            self.handle_message(message)?;
+        }
+
+        Ok(())
+    }
+
+    /// This is a convenience method that's same as [`receive`], but panics if the server got
+    /// disconnected. It will pass other errors as is.
+    ///
+    /// [`receive`]: TestServer::receive
+    fn receive_or_panic(&mut self) -> Result<(), TestServerError> {
+        if let Err(err) = self.receive(None) {
+            if err.is_disconnected() {
+                self.panic_on_server_disconnect();
+            } else {
+                return Err(err);
+            }
+        }
+        Ok(())
+    }
+
+    /// Handle the incoming message from the server.
+    ///
+    /// This method will store the message as follows:
+    /// - Requests are stored in `self.requests`
+    /// - Responses are stored in `self.responses` with the request ID as the key
+    /// - Notifications are stored in `self.notifications`
+    fn handle_message(&mut self, message: Message) -> Result<(), TestServerError> {
+        match message {
+            Message::Request(request) => {
+                self.requests.push_back(request);
+            }
+            Message::Response(response) => match self.responses.entry(response.id.clone()) {
+                Entry::Occupied(existing) => {
+                    return Err(TestServerError::DuplicateResponse(
+                        response.id,
+                        Box::new(existing.get().clone()),
+                    ));
+                }
+                Entry::Vacant(entry) => {
+                    entry.insert(response);
+                }
+            },
+            Message::Notification(notification) => {
+                self.notifications.push_back(notification);
+            }
+        }
+        Ok(())
+    }
+
+    #[track_caller]
+    fn panic_on_server_disconnect(&mut self) -> ! {
+        if let Some(handle) = &self.server_thread {
+            if handle.is_finished() {
+                let handle = self.server_thread.take().unwrap();
+                if let Err(panic) = handle.join() {
+                    std::panic::resume_unwind(panic);
+                }
+                panic!("Server exited unexpectedly");
+            }
+        }
+
+        panic!("Server dropped client receiver while still running");
+    }
+
+    /// Handle workspace configuration requests from the server.
+    ///
+    /// Use the [`get_request`] method to wait for the server to send this request.
+    ///
+    /// [`get_request`]: TestServer::get_request
+    pub(crate) fn handle_workspace_configuration_request(
+        &mut self,
+        request_id: RequestId,
+        params: &ConfigurationParams,
+    ) -> Result<()> {
+        let mut results = Vec::new();
+
+        for item in &params.items {
+            let Some(scope_uri) = &item.scope_uri else {
+                unimplemented!("Handling global configuration requests is not implemented yet");
+            };
+            let config_value = if let Some(options) = self.workspace_configurations.get(scope_uri) {
+                // Return the configuration for the specific workspace
+                match item.section.as_deref() {
+                    Some("ty") => serde_json::to_value(options)?,
+                    Some(_) | None => {
+                        // TODO: Handle `python` section once it's implemented in the server
+                        // As per the spec:
+                        //
+                        // > If the client can't provide a configuration setting for a given scope
+                        // > then null needs to be present in the returned array.
+                        serde_json::Value::Null
+                    }
+                }
+            } else {
+                tracing::warn!("No workspace configuration found for {scope_uri}");
+                serde_json::Value::Null
+            };
+            results.push(config_value);
+        }
+
+        let response = Response::new_ok(request_id, results);
+        self.send(Message::Response(response));
+
+        Ok(())
+    }
+
+    /// Get the initialization result
+    pub(crate) fn initialization_result(&self) -> Option<&InitializeResult> {
+        self.initialize_response.as_ref()
+    }
+
+    fn file_uri(&self, path: impl AsRef<SystemPath>) -> Url {
+        Url::from_file_path(self.test_dir.root().join(path.as_ref()).as_std_path())
+            .expect("Path must be a valid URL")
+    }
+
+    /// Send a `textDocument/didOpen` notification
+    pub(crate) fn open_text_document(
+        &mut self,
+        path: impl AsRef<SystemPath>,
+        content: &impl ToString,
+        version: i32,
+    ) {
+        let params = DidOpenTextDocumentParams {
+            text_document: TextDocumentItem {
+                uri: self.file_uri(path),
+                language_id: "python".to_string(),
+                version,
+                text: content.to_string(),
+            },
+        };
+        self.send_notification::<DidOpenTextDocument>(params);
+    }
+
+    /// Send a `textDocument/didChange` notification with the given content changes
+    #[expect(dead_code)]
+    pub(crate) fn change_text_document(
+        &mut self,
+        path: impl AsRef<SystemPath>,
+        changes: Vec<TextDocumentContentChangeEvent>,
+        version: i32,
+    ) {
+        let params = DidChangeTextDocumentParams {
+            text_document: VersionedTextDocumentIdentifier {
+                uri: self.file_uri(path),
+                version,
+            },
+            content_changes: changes,
+        };
+        self.send_notification::<DidChangeTextDocument>(params);
+    }
+
+    /// Send a `textDocument/didClose` notification
+    #[expect(dead_code)]
+    pub(crate) fn close_text_document(&mut self, path: impl AsRef<SystemPath>) {
+        let params = DidCloseTextDocumentParams {
+            text_document: TextDocumentIdentifier {
+                uri: self.file_uri(path),
+            },
+        };
+        self.send_notification::<DidCloseTextDocument>(params);
+    }
+
+    /// Send a `workspace/didChangeWatchedFiles` notification with the given file events
+    #[expect(dead_code)]
+    pub(crate) fn did_change_watched_files(&mut self, events: Vec<FileEvent>) {
+        let params = DidChangeWatchedFilesParams { changes: events };
+        self.send_notification::<DidChangeWatchedFiles>(params);
+    }
+
+    /// Send a `textDocument/diagnostic` request for the document at the given path.
+    pub(crate) fn document_diagnostic_request(
+        &mut self,
+        path: impl AsRef<SystemPath>,
+    ) -> Result<DocumentDiagnosticReportResult> {
+        let params = DocumentDiagnosticParams {
+            text_document: TextDocumentIdentifier {
+                uri: self.file_uri(path),
+            },
+            identifier: Some("ty".to_string()),
+            previous_result_id: None,
+            work_done_progress_params: WorkDoneProgressParams::default(),
+            partial_result_params: PartialResultParams::default(),
+        };
+        let id = self.send_request::<DocumentDiagnosticRequest>(params);
+        self.await_response::<DocumentDiagnosticReportResult>(id)
+    }
+}
+
+impl fmt::Debug for TestServer {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("TestServer")
+            .field("temp_dir", &self.test_dir.root())
+            .field("request_counter", &self.request_counter)
+            .field("responses", &self.responses)
+            .field("notifications", &self.notifications)
+            .field("server_requests", &self.requests)
+            .field("initialize_response", &self.initialize_response)
+            .field("workspace_configurations", &self.workspace_configurations)
+            .field("registered_capabilities", &self.registered_capabilities)
+            .finish_non_exhaustive()
+    }
+}
+
+impl Drop for TestServer {
+    fn drop(&mut self) {
+        self.drain_messages();
+
+        // Follow the LSP protocol to shutdown the server gracefully.
+        //
+        // The `server_thread` could be `None` if the server exited unexpectedly or panicked or if
+        // it dropped the client connection.
+        let shutdown_error = if self.server_thread.is_some() {
+            let shutdown_id = self.send_request::<Shutdown>(());
+            match self.await_response::<()>(shutdown_id) {
+                Ok(()) => {
+                    self.send_notification::<Exit>(());
+                    None
+                }
+                Err(err) => Some(format!("Failed to get shutdown response: {err:?}")),
+            }
+        } else {
+            None
+        };
+
+        if let Some(_client_connection) = self.client_connection.take() {
+            // Drop the client connection before joining the server thread to avoid any hangs
+            // in case the server didn't respond to the shutdown request.
+        }
+
+        if std::thread::panicking() {
+            // If the test server panicked, avoid further assertions.
+            return;
+        }
+
+        if let Some(server_thread) = self.server_thread.take() {
+            if let Err(err) = server_thread.join() {
+                panic!("Panic in the server thread: {err:?}");
+            }
+        }
+
+        if let Some(error) = shutdown_error {
+            panic!("Test server did not shut down gracefully: {error}");
+        }
+
+        self.assert_no_pending_messages();
+    }
+}
+
+/// Builder for creating test servers with specific configurations
+pub(crate) struct TestServerBuilder {
+    test_dir: TestContext,
+    workspaces: Vec<(WorkspaceFolder, ClientOptions)>,
+    client_capabilities: ClientCapabilities,
+}
+
+impl TestServerBuilder {
+    /// Create a new builder
+    pub(crate) fn new() -> Result<Self> {
+        // Default client capabilities for the test server. These are assumptions made by the real
+        // server and are common for most clients:
+        //
+        // - Supports publishing diagnostics
+        // - Supports pulling workspace configuration
+        let client_capabilities = ClientCapabilities {
+            text_document: Some(TextDocumentClientCapabilities {
+                publish_diagnostics: Some(PublishDiagnosticsClientCapabilities::default()),
+                ..Default::default()
+            }),
+            workspace: Some(WorkspaceClientCapabilities {
+                configuration: Some(true),
+                ..Default::default()
+            }),
+            ..Default::default()
+        };
+
+        Ok(Self {
+            workspaces: Vec::new(),
+            test_dir: TestContext::new()?,
+            client_capabilities,
+        })
+    }
+
+    /// Add a workspace to the test server with the given root path and options.
+    ///
+    /// This option will be used to respond to the `workspace/configuration` request that the
+    /// server will send to the client.
+    pub(crate) fn with_workspace(
+        mut self,
+        workspace_root: &SystemPath,
+        options: ClientOptions,
+    ) -> Result<Self> {
+        // TODO: Support multiple workspaces in the test server
+        if self.workspaces.len() == 1 {
+            anyhow::bail!("Test server doesn't support multiple workspaces yet");
+        }
+
+        let workspace_path = self.test_dir.root().join(workspace_root);
+        fs::create_dir_all(workspace_path.as_std_path())?;
+
+        self.workspaces.push((
+            WorkspaceFolder {
+                uri: Url::from_file_path(workspace_path.as_std_path()).map_err(|()| {
+                    anyhow!("Failed to convert workspace path to URL: {workspace_path}")
+                })?,
+                name: workspace_root.file_name().unwrap_or("test").to_string(),
+            },
+            options,
+        ));
+
+        Ok(self)
+    }
+
+    /// Enable or disable pull diagnostics capability
+    pub(crate) fn enable_pull_diagnostics(mut self, enabled: bool) -> Self {
+        self.client_capabilities
+            .text_document
+            .get_or_insert_with(Default::default)
+            .diagnostic = if enabled {
+            Some(DiagnosticClientCapabilities::default())
+        } else {
+            None
+        };
+        self
+    }
+
+    /// Enable or disable file watching capability
+    #[expect(dead_code)]
+    pub(crate) fn enable_did_change_watched_files(mut self, enabled: bool) -> Self {
+        self.client_capabilities
+            .workspace
+            .get_or_insert_with(Default::default)
+            .did_change_watched_files = if enabled {
+            Some(DidChangeWatchedFilesClientCapabilities::default())
+        } else {
+            None
+        };
+        self
+    }
+
+    /// Set custom client capabilities (overrides any previously set capabilities)
+    #[expect(dead_code)]
+    pub(crate) fn with_client_capabilities(mut self, capabilities: ClientCapabilities) -> Self {
+        self.client_capabilities = capabilities;
+        self
+    }
+
+    /// Write a file to the test directory
+    pub(crate) fn with_file(
+        self,
+        path: impl AsRef<SystemPath>,
+        content: impl AsRef<str>,
+    ) -> Result<Self> {
+        let file_path = self.test_dir.root().join(path.as_ref());
+        // Ensure parent directories exists
+        if let Some(parent) = file_path.parent() {
+            fs::create_dir_all(parent.as_std_path())?;
+        }
+        fs::write(file_path.as_std_path(), content.as_ref())?;
+        Ok(self)
+    }
+
+    /// Write multiple files to the temporary directory
+    #[expect(dead_code)]
+    pub(crate) fn with_files<P, C, I>(mut self, files: I) -> Result<Self>
+    where
+        I: IntoIterator<Item = (P, C)>,
+        P: AsRef<SystemPath>,
+        C: AsRef<str>,
+    {
+        for (path, content) in files {
+            self = self.with_file(path, content)?;
+        }
+        Ok(self)
+    }
+
+    /// Build the test server
+    pub(crate) fn build(self) -> Result<TestServer> {
+        TestServer::new(self.workspaces, self.test_dir, self.client_capabilities)
+    }
+}
+
+/// A context specific to a server test.
+///
+/// This creates a temporary directory that is used as the current working directory for the server
+/// in which the test files are stored. This also holds the insta settings scope that filters out
+/// the temporary directory path from snapshots.
+///
+/// This is similar to the `CliTest` in `ty` crate.
+struct TestContext {
+    _temp_dir: TempDir,
+    _settings_scope: SettingsBindDropGuard,
+    project_dir: SystemPathBuf,
+}
+
+impl TestContext {
+    pub(crate) fn new() -> anyhow::Result<Self> {
+        let temp_dir = TempDir::new()?;
+
+        // Canonicalize the tempdir path because macos uses symlinks for tempdirs
+        // and that doesn't play well with our snapshot filtering.
+        // Simplify with dunce because otherwise we get UNC paths on Windows.
+        let project_dir = SystemPathBuf::from_path_buf(
+            dunce::simplified(
+                &temp_dir
+                    .path()
+                    .canonicalize()
+                    .context("Failed to canonicalize project path")?,
+            )
+            .to_path_buf(),
+        )
+        .map_err(|path| {
+            anyhow!(
+                "Failed to create test directory: `{}` contains non-Unicode characters",
+                path.display()
+            )
+        })?;
+
+        let mut settings = insta::Settings::clone_current();
+        settings.add_filter(&tempdir_filter(project_dir.as_str()), "<temp_dir>/");
+        settings.add_filter(
+            &tempdir_filter(
+                Url::from_file_path(project_dir.as_std_path())
+                    .map_err(|()| anyhow!("Failed to convert root directory to url"))?
+                    .path(),
+            ),
+            "<temp_dir>/",
+        );
+        settings.add_filter(r#"\\(\w\w|\s|\.|")"#, "/$1");
+        settings.add_filter(
+            r#"The system cannot find the file specified."#,
+            "No such file or directory",
+        );
+
+        let settings_scope = settings.bind_to_scope();
+
+        Ok(Self {
+            project_dir,
+            _temp_dir: temp_dir,
+            _settings_scope: settings_scope,
+        })
+    }
+
+    pub(crate) fn root(&self) -> &SystemPath {
+        &self.project_dir
+    }
+}
+
+fn tempdir_filter(path: impl AsRef<str>) -> String {
+    format!(r"{}\\?/?", regex::escape(path.as_ref()))
+}