Add test case for missing VIRTUAL_ENV

crates/ruff_db/src/system/test.rs

@@ -1,5 +1,6 @@
 use glob::PatternError;
 use ruff_notebook::{Notebook, NotebookError};
+use rustc_hash::FxHashMap;
 use std::panic::RefUnwindSafe;
 use std::sync::{Arc, Mutex};
 
@@ -20,18 +21,47 @@ use super::walk_directory::WalkDirectoryBuilder;
 ///
 /// ## Warning
 /// Don't use this system for production code. It's intended for testing only.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct TestSystem {
     inner: Arc<dyn WritableSystem + RefUnwindSafe + Send + Sync>,
+    /// Environment variable overrides. If a key is present here, it takes precedence
+    /// over the inner system's environment variables.
+    env_overrides: Arc<Mutex<FxHashMap<String, Option<String>>>>,
+}
+
+impl Clone for TestSystem {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+            env_overrides: self.env_overrides.clone(),
+        }
+    }
 }
 
 impl TestSystem {
     pub fn new(inner: impl WritableSystem + RefUnwindSafe + Send + Sync + 'static) -> Self {
         Self {
             inner: Arc::new(inner),
+            env_overrides: Arc::new(Mutex::new(FxHashMap::default())),
         }
     }
 
+    /// Sets an environment variable override. This takes precedence over the inner system.
+    pub fn set_env_var(&self, name: impl Into<String>, value: impl Into<String>) {
+        self.env_overrides
+            .lock()
+            .unwrap()
+            .insert(name.into(), Some(value.into()));
+    }
+
+    /// Removes an environment variable override, making it appear as not set.
+    pub fn remove_env_var(&self, name: impl Into<String>) {
+        self.env_overrides
+            .lock()
+            .unwrap()
+            .insert(name.into(), None);
+    }
+
     /// Returns the [`InMemorySystem`].
     ///
     /// ## Panics
@@ -147,6 +177,18 @@ impl System for TestSystem {
         self.system().case_sensitivity()
     }
 
+    fn env_var(&self, name: &str) -> std::result::Result<String, std::env::VarError> {
+        // Check overrides first
+        if let Some(override_value) = self.env_overrides.lock().unwrap().get(name) {
+            return match override_value {
+                Some(value) => Ok(value.clone()),
+                None => Err(std::env::VarError::NotPresent),
+            };
+        }
+        // Fall back to inner system
+        self.system().env_var(name)
+    }
+
     fn dyn_clone(&self) -> Box<dyn System> {
         Box::new(self.clone())
     }
@@ -156,6 +198,7 @@ impl Default for TestSystem {
     fn default() -> Self {
         Self {
             inner: Arc::new(InMemorySystem::default()),
+            env_overrides: Arc::new(Mutex::new(FxHashMap::default())),
         }
     }
 }

crates/ty/tests/cli/python_environment.rs

@@ -2703,3 +2703,51 @@ fn pythonpath_multiple_dirs_is_respected() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+/// Test behavior when `VIRTUAL_ENV` is set but points to a non-existent path.
+#[test]
+fn missing_virtual_env() -> anyhow::Result<()> {
+    let working_venv_package1_path = if cfg!(windows) {
+        "project/.venv/Lib/site-packages/package1/__init__.py"
+    } else {
+        "project/.venv/lib/python3.13/site-packages/package1/__init__.py"
+    };
+
+    let case = CliTest::with_files([
+        (
+            "project/test.py",
+            r#"
+            from package1 import WorkingVenv
+            "#,
+        ),
+        (
+            "project/.venv/pyvenv.cfg",
+            r#"
+home = ./
+
+            "#,
+        ),
+        (
+            working_venv_package1_path,
+            r#"
+            class WorkingVenv: ...
+            "#,
+        ),
+    ])?;
+
+    assert_cmd_snapshot!(case.command()
+        .current_dir(case.root().join("project"))
+        .env("VIRTUAL_ENV", case.root().join("nonexistent-venv")), @r"
+    success: false
+    exit_code: 2
+    ----- stdout -----
+
+    ----- stderr -----
+    ty failed
+      Cause: Failed to discover local Python environment
+      Cause: Invalid `VIRTUAL_ENV` environment variable `<temp_dir>/nonexistent-venv`: does not point to a directory on disk
+      Cause: No such file or directory (os error 2)
+    ");
+
+    Ok(())
+}

crates/ty_server/tests/e2e/initialize.rs

@@ -1,5 +1,6 @@
 use anyhow::Result;
-use lsp_types::{Position, notification::ShowMessage, request::RegisterCapability};
+use lsp_types::notification::ShowMessage;
+use lsp_types::{Position, request::RegisterCapability};
 use ruff_db::system::SystemPath;
 use serde_json::Value;
 use ty_server::{ClientOptions, DiagnosticMode};
@@ -474,3 +475,20 @@ fn register_multiple_capabilities() -> Result<()> {
 
     Ok(())
 }
+
+/// Tests that the server doesn't panic when `VIRTUAL_ENV` points to a non-existent directory.
+///
+/// See: <https://github.com/astral-sh/ty/issues/2031>
+#[test]
+fn missing_virtual_env_does_not_panic() -> Result<()> {
+    let workspace_root = SystemPath::new("project");
+
+    // This should not panic even though VIRTUAL_ENV points to a non-existent path
+    let _server = TestServerBuilder::new()?
+        .with_workspace(workspace_root, None)?
+        .with_env_var("VIRTUAL_ENV", "/nonexistent/virtual/env/path")
+        .build()
+        .wait_until_workspaces_are_initialized();
+
+    Ok(())
+}

crates/ty_server/tests/e2e/main.rs

@@ -209,6 +209,7 @@ impl TestServer {
         test_context: TestContext,
         capabilities: ClientCapabilities,
         initialization_options: Option<ClientOptions>,
+        env_vars: Vec<(String, String)>,
     ) -> Self {
         setup_tracing();
 
@@ -219,11 +220,16 @@ impl TestServer {
         // Create OS system with the test directory as cwd
         let os_system = OsSystem::new(test_context.root());
 
+        // Create test system and set environment variable overrides
+        let test_system = Arc::new(TestSystem::new(os_system));
+        for (name, value) in env_vars {
+            test_system.set_env_var(name, value);
+        }
+
         // 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, true) {
                 Ok(server) => {
@@ -1052,6 +1058,7 @@ pub(crate) struct TestServerBuilder {
     workspaces: Vec<(WorkspaceFolder, Option<ClientOptions>)>,
     initialization_options: Option<ClientOptions>,
     client_capabilities: ClientCapabilities,
+    env_vars: Vec<(String, String)>,
 }
 
 impl TestServerBuilder {
@@ -1082,6 +1089,7 @@ impl TestServerBuilder {
             test_context: TestContext::new()?,
             initialization_options: None,
             client_capabilities,
+            env_vars: Vec::new(),
         })
     }
 
@@ -1091,6 +1099,16 @@ impl TestServerBuilder {
         self
     }
 
+    /// Set an environment variable for the test server's system.
+    pub(crate) fn with_env_var(
+        mut self,
+        name: impl Into<String>,
+        value: impl Into<String>,
+    ) -> Self {
+        self.env_vars.push((name.into(), value.into()));
+        self
+    }
+
     /// 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
@@ -1237,6 +1255,7 @@ impl TestServerBuilder {
             self.test_context,
             self.client_capabilities,
             self.initialization_options,
+            self.env_vars,
         )
     }
 }