mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-10-02 14:51:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

[ty] Move server tests as integration tests (#19522)

Browse source 
## Summary

Reference:
https://github.com/astral-sh/ruff/pull/19391#discussion_r2222780892
This commit is contained in:
Dhruv Manilawala 2025-07-24 21:40:17 +05:30 • committed by GitHub
parent 1d2181623c
commit f9091ea8bb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 172 additions and 128 deletions
Download patch file Download diff file Expand all files Collapse all files

7
crates/ty_server/src/lib.rs
View file

@ -4,7 +4,9 @@ use anyhow::Context;
use lsp_server::Connection;
use ruff_db::system::{OsSystem, SystemPathBuf};
use crate::server::Server;
pub use crate::logging::{LogLevel, init_logging};
pub use crate::server::Server;
pub use crate::session::ClientOptions;
pub use document::{NotebookDocument, PositionEncoding, TextDocument};
pub(crate) use session::{DocumentQuery, Session};
@ -14,9 +16,6 @@ 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";

4
crates/ty_server/src/logging.rs
View file

@ -14,7 +14,7 @@ use tracing_subscriber::fmt::time::ChronoLocal;
use tracing_subscriber::fmt::writer::BoxMakeWriter;
use tracing_subscriber::layer::SubscriberExt;
pub(crate) fn init_logging(log_level: LogLevel, log_file: Option<&SystemPath>) {
pub fn init_logging(log_level: LogLevel, log_file: Option<&SystemPath>) {
let log_file = log_file
.map(|path| {
// this expands `logFile` so that tildes and environment variables
@ -66,7 +66,7 @@ pub(crate) fn init_logging(log_level: LogLevel, log_file: Option<&SystemPath>) {
/// The default log level is `info`.
#[derive(Clone, Copy, Debug, Deserialize, Default, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "lowercase")]
pub(crate) enum LogLevel {
pub enum LogLevel {
Error,
Warn,
#[default]

92
crates/ty_server/src/server.rs
View file

@ -26,7 +26,7 @@ pub(crate) use api::publish_settings_diagnostics;
pub(crate) use main_loop::{Action, ConnectionSender, Event, MainLoopReceiver, MainLoopSender};
pub(crate) type Result<T> = std::result::Result<T, api::Error>;
pub(crate) struct Server {
pub struct Server {
connection: Connection,
client_capabilities: ClientCapabilities,
worker_threads: NonZeroUsize,
@ -36,7 +36,7 @@ pub(crate) struct Server {
}
impl Server {
pub(crate) fn new(
pub fn new(
worker_threads: NonZeroUsize,
connection: Connection,
native_system: Arc<dyn System + 'static + Send + Sync + RefUnwindSafe>,
@ -161,7 +161,7 @@ impl Server {
})
}
pub(crate) fn run(mut self) -> crate::Result<()> {
pub fn run(mut self) -> crate::Result<()> {
let client = Client::new(
self.main_loop_sender.clone(),
self.connection.sender.clone(),
@ -302,89 +302,3 @@ 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(())
}
}

3
crates/ty_server/src/session.rs
View file

@ -21,7 +21,8 @@ use ty_project::{ChangeResult, Db as _, ProjectDatabase, ProjectMetadata};
pub(crate) use self::capabilities::ResolvedClientCapabilities;
pub(crate) use self::index::DocumentQuery;
pub(crate) use self::options::{AllOptions, ClientOptions, DiagnosticMode};
pub use self::options::ClientOptions;
pub(crate) use self::options::{AllOptions, DiagnosticMode};
pub(crate) use self::settings::ClientSettings;
use crate::document::{DocumentKey, DocumentVersion, NotebookDocument};
use crate::server::publish_settings_diagnostics;

29
crates/ty_server/src/session/options.rs
View file

@ -49,9 +49,10 @@ struct WorkspaceOptions {
/// This is a direct representation of the settings schema sent by the client.
#[derive(Clone, Debug, Deserialize, Default)]
#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[cfg_attr(feature = "testing", derive(serde::Serialize))]
#[serde(rename_all = "camelCase")]
pub(crate) struct ClientOptions {
pub struct ClientOptions {
/// Settings under the `python.*` namespace in VS Code that are useful for the ty language
/// server.
python: Option<Python>,
@ -63,7 +64,8 @@ pub(crate) struct ClientOptions {
/// Diagnostic mode for the language server.
#[derive(Clone, Copy, Debug, Default, Deserialize)]
#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[cfg_attr(feature = "testing", derive(serde::Serialize))]
#[serde(rename_all = "camelCase")]
pub(crate) enum DiagnosticMode {
/// Check only currently open files.
@ -147,21 +149,24 @@ impl ClientOptions {
// all settings and not just the ones in "python.*".
#[derive(Clone, Debug, Deserialize, Default)]
#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[cfg_attr(feature = "testing", derive(serde::Serialize))]
#[serde(rename_all = "camelCase")]
struct Python {
ty: Option<Ty>,
}
#[derive(Clone, Debug, Deserialize, Default)]
#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[cfg_attr(feature = "testing", derive(serde::Serialize))]
#[serde(rename_all = "camelCase")]
struct PythonExtension {
active_environment: Option<ActiveEnvironment>,
}
#[derive(Clone, Debug, Deserialize)]
#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[cfg_attr(feature = "testing", derive(serde::Serialize))]
#[serde(rename_all = "camelCase")]
pub(crate) struct ActiveEnvironment {
pub(crate) executable: PythonExecutable,
@ -170,7 +175,8 @@ pub(crate) struct ActiveEnvironment {
}
#[derive(Clone, Debug, Deserialize)]
#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[cfg_attr(feature = "testing", derive(serde::Serialize))]
#[serde(rename_all = "camelCase")]
pub(crate) struct EnvironmentVersion {
pub(crate) major: i64,
@ -182,7 +188,8 @@ pub(crate) struct EnvironmentVersion {
}
#[derive(Clone, Debug, Deserialize)]
#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[cfg_attr(feature = "testing", derive(serde::Serialize))]
#[serde(rename_all = "camelCase")]
pub(crate) struct PythonEnvironment {
pub(crate) folder_uri: Url,
@ -194,7 +201,8 @@ pub(crate) struct PythonEnvironment {
}
#[derive(Clone, Debug, Deserialize)]
#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[cfg_attr(feature = "testing", derive(serde::Serialize))]
#[serde(rename_all = "camelCase")]
pub(crate) struct PythonExecutable {
#[allow(dead_code)]
@ -203,7 +211,8 @@ pub(crate) struct PythonExecutable {
}
#[derive(Clone, Debug, Deserialize, Default)]
#[cfg_attr(test, derive(serde::Serialize, PartialEq, Eq))]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[cfg_attr(feature = "testing", derive(serde::Serialize))]
#[serde(rename_all = "camelCase")]
struct Ty {
disable_language_services: Option<bool>,

70
crates/ty_server/src/snapshots/ty_server__server__tests__initialization.snap
View file

@ -1,70 +0,0 @@
---
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,
"referencesProvider": 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"
}
}

70
crates/ty_server/src/snapshots/ty_server__server__tests__initialization_with_workspace.snap
View file

@ -1,70 +0,0 @@
---
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,
"referencesProvider": 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"
}
}

99
crates/ty_server/src/snapshots/ty_server__server__tests__publish_diagnostics_on_did_open.snap
View file

@ -1,99 +0,0 @@
---
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,
),
}

93
crates/ty_server/src/snapshots/ty_server__server__tests__pull_diagnostics_on_did_open.snap
View file

@ -1,93 +0,0 @@
---
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,
},
],
},
},
),
)

889
crates/ty_server/src/test.rs
View file

@ -1,889 +0,0 @@
//! 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()))
}