ruff_db/system/

test.rs

use ruff_notebook::{Notebook, NotebookError};
use rustc_hash::FxHashMap;
use std::panic::RefUnwindSafe;
use std::sync::{Arc, Mutex};

use crate::Db;
use crate::files::File;
use crate::system::{
    CommandExecutor, DirectoryEntry, MemoryFileSystem, Metadata, Result, System, SystemPath,
    SystemPathBuf, SystemVirtualPath, WhichError, WhichResult,
};

use super::WritableSystem;
use super::walk_directory::WalkDirectoryBuilder;

/// System implementation intended for testing.
///
/// It uses a memory-file system by default, but can be switched to the real file system for tests
/// verifying more advanced file system features.
///
/// ## Warning
/// Don't use this system for production code. It's intended for testing only.
#[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
    /// If the underlying test system isn't the [`InMemorySystem`].
    pub fn in_memory(&self) -> &InMemorySystem {
        self.as_in_memory()
            .expect("The test db is not using a memory file system")
    }

    /// Returns the `InMemorySystem` or `None` if the underlying test system isn't the [`InMemorySystem`].
    fn as_in_memory(&self) -> Option<&InMemorySystem> {
        self.system().as_any().downcast_ref::<InMemorySystem>()
    }

    /// Returns the memory file system.
    ///
    /// ## Panics
    /// If the underlying test system isn't the [`InMemorySystem`].
    pub fn memory_file_system(&self) -> &MemoryFileSystem {
        self.in_memory().fs()
    }

    fn use_system<S>(&mut self, system: S)
    where
        S: WritableSystem + Send + Sync + RefUnwindSafe + 'static,
    {
        self.inner = Arc::new(system);
    }

    fn system(&self) -> &dyn WritableSystem {
        &*self.inner
    }
}

impl System for TestSystem {
    fn path_metadata(&self, path: &SystemPath) -> Result<Metadata> {
        self.system().path_metadata(path)
    }

    fn canonicalize_path(&self, path: &SystemPath) -> Result<SystemPathBuf> {
        self.system().canonicalize_path(path)
    }

    fn is_same_file(&self, first: &SystemPath, second: &SystemPath) -> Result<bool> {
        self.system().is_same_file(first, second)
    }

    fn read_to_string(&self, path: &SystemPath) -> Result<String> {
        self.system().read_to_string(path)
    }

    fn read_to_notebook(&self, path: &SystemPath) -> std::result::Result<Notebook, NotebookError> {
        self.system().read_to_notebook(path)
    }

    fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result<String> {
        self.system().read_virtual_path_to_string(path)
    }

    fn read_virtual_path_to_notebook(
        &self,
        path: &SystemVirtualPath,
    ) -> std::result::Result<Notebook, NotebookError> {
        self.system().read_virtual_path_to_notebook(path)
    }

    fn current_directory(&self) -> &SystemPath {
        self.system().current_directory()
    }

    fn user_config_directory(&self) -> Option<SystemPathBuf> {
        self.system().user_config_directory()
    }

    fn cache_dir(&self) -> Option<SystemPathBuf> {
        self.system().cache_dir()
    }

    fn which(&self, _name: &str) -> WhichResult {
        Err(WhichError::CannotFindBinaryPath)
    }

    fn command_executor(&self) -> Option<&dyn CommandExecutor> {
        self.system().command_executor()
    }

    fn read_directory<'a>(
        &'a self,
        path: &SystemPath,
    ) -> Result<Box<dyn Iterator<Item = Result<DirectoryEntry>> + 'a>> {
        self.system().read_directory(path)
    }

    fn walk_directory(&self, path: &SystemPath) -> WalkDirectoryBuilder {
        self.system().walk_directory(path)
    }

    fn as_writable(&self) -> Option<&dyn WritableSystem> {
        Some(self)
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
        self
    }

    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())
    }
}

impl Default for TestSystem {
    fn default() -> Self {
        Self::new(InMemorySystem::default())
    }
}

impl WritableSystem for TestSystem {
    fn create_new_file(&self, path: &SystemPath) -> Result<()> {
        self.system().create_new_file(path)
    }

    fn write_file_bytes(&self, path: &SystemPath, content: &[u8]) -> Result<()> {
        self.system().write_file_bytes(path, content)
    }

    fn create_directory_all(&self, path: &SystemPath) -> Result<()> {
        self.system().create_directory_all(path)
    }

    fn dyn_clone(&self) -> Box<dyn WritableSystem> {
        Box::new(self.clone())
    }
}

/// Extension trait for databases that use a [`WritableSystem`].
///
/// Provides various helper function that ease testing.
pub trait DbWithWritableSystem: Db + Sized {
    type System: WritableSystem;

    fn writable_system(&self) -> &Self::System;

    /// Writes the content of the given file and notifies the Db about the change.
    fn write_file(&mut self, path: impl AsRef<SystemPath>, content: impl AsRef<str>) -> Result<()> {
        let path = path.as_ref();
        match self.writable_system().write_file(path, content.as_ref()) {
            Ok(()) => {
                File::sync_path(self, path);
                Ok(())
            }
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
                if let Some(parent) = path.parent() {
                    self.writable_system().create_directory_all(parent)?;

                    for ancestor in parent.ancestors() {
                        File::sync_path(self, ancestor);
                    }

                    self.writable_system().write_file(path, content.as_ref())?;
                    File::sync_path(self, path);

                    Ok(())
                } else {
                    Err(error)
                }
            }
            err => err,
        }
    }

    /// Writes auto-dedented text to a file.
    fn write_dedented(&mut self, path: &str, content: &str) -> Result<()> {
        self.write_file(path, ruff_python_trivia::textwrap::dedent(content))?;
        Ok(())
    }

    /// Writes the content of the given files and notifies the Db about the change.
    fn write_files<P, C, I>(&mut self, files: I) -> Result<()>
    where
        I: IntoIterator<Item = (P, C)>,
        P: AsRef<SystemPath>,
        C: AsRef<str>,
    {
        for (path, content) in files {
            self.write_file(path, content)?;
        }

        Ok(())
    }
}

/// Extension trait for databases that use [`TestSystem`].
///
/// Provides various helper function that ease testing.
pub trait DbWithTestSystem: Db + Sized {
    fn test_system(&self) -> &TestSystem;

    fn test_system_mut(&mut self) -> &mut TestSystem;

    /// Writes the content of the given virtual file.
    ///
    /// ## Panics
    /// If the db isn't using the [`InMemorySystem`].
    fn write_virtual_file(
        &mut self,
        path: impl AsRef<SystemVirtualPath>,
        content: impl AsRef<[u8]>,
    ) {
        let path = path.as_ref();
        self.test_system()
            .memory_file_system()
            .write_virtual_file(path, content);
    }

    /// Uses the given system instead of the testing system.
    ///
    /// This useful for testing advanced file system features like permissions, symlinks, etc.
    ///
    /// Note that any files written to the memory file system won't be copied over.
    fn use_system<S>(&mut self, os: S)
    where
        S: WritableSystem + Send + Sync + RefUnwindSafe + 'static,
    {
        self.test_system_mut().use_system(os);
    }

    /// Returns the memory file system.
    ///
    /// ## Panics
    /// If the underlying test system isn't the [`InMemorySystem`].
    fn memory_file_system(&self) -> &MemoryFileSystem {
        self.test_system().memory_file_system()
    }
}

impl<T> DbWithWritableSystem for T
where
    T: DbWithTestSystem,
{
    type System = TestSystem;

    fn writable_system(&self) -> &Self::System {
        self.test_system()
    }
}

#[derive(Clone, Default, Debug)]
pub struct InMemorySystem {
    user_config_directory: Arc<Mutex<Option<SystemPathBuf>>>,
    memory_fs: MemoryFileSystem,
}

impl InMemorySystem {
    pub fn from_memory_fs(memory_fs: MemoryFileSystem) -> Self {
        Self {
            user_config_directory: Mutex::new(None).into(),
            memory_fs,
        }
    }

    pub fn fs(&self) -> &MemoryFileSystem {
        &self.memory_fs
    }

    pub fn set_user_configuration_directory(&self, directory: Option<SystemPathBuf>) {
        let mut user_directory = self.user_config_directory.lock().unwrap();
        *user_directory = directory;
    }
}

impl System for InMemorySystem {
    fn path_metadata(&self, path: &SystemPath) -> Result<Metadata> {
        self.memory_fs.metadata(path)
    }

    fn canonicalize_path(&self, path: &SystemPath) -> Result<SystemPathBuf> {
        self.memory_fs.canonicalize(path)
    }

    fn is_same_file(&self, first: &SystemPath, second: &SystemPath) -> Result<bool> {
        // The in-memory file system does not support hard links, so canonical paths uniquely
        // identify files.
        Ok(self.canonicalize_path(first)? == self.canonicalize_path(second)?)
    }

    fn read_to_string(&self, path: &SystemPath) -> Result<String> {
        self.memory_fs.read_to_string(path)
    }

    fn read_to_notebook(&self, path: &SystemPath) -> std::result::Result<Notebook, NotebookError> {
        let content = self.read_to_string(path)?;
        Notebook::from_source_code(&content)
    }

    fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result<String> {
        self.memory_fs.read_virtual_path_to_string(path)
    }

    fn read_virtual_path_to_notebook(
        &self,
        path: &SystemVirtualPath,
    ) -> std::result::Result<Notebook, NotebookError> {
        let content = self.read_virtual_path_to_string(path)?;
        Notebook::from_source_code(&content)
    }

    fn current_directory(&self) -> &SystemPath {
        self.memory_fs.current_directory()
    }

    fn user_config_directory(&self) -> Option<SystemPathBuf> {
        self.user_config_directory.lock().unwrap().clone()
    }

    fn cache_dir(&self) -> Option<SystemPathBuf> {
        None
    }

    fn which(&self, _name: &str) -> WhichResult {
        Err(WhichError::CannotFindBinaryPath)
    }

    fn read_directory<'a>(
        &'a self,
        path: &SystemPath,
    ) -> Result<Box<dyn Iterator<Item = Result<DirectoryEntry>> + 'a>> {
        Ok(Box::new(self.memory_fs.read_directory(path)?))
    }

    fn walk_directory(&self, path: &SystemPath) -> WalkDirectoryBuilder {
        self.memory_fs.walk_directory(path)
    }

    fn as_writable(&self) -> Option<&dyn WritableSystem> {
        Some(self)
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
        self
    }

    fn dyn_clone(&self) -> Box<dyn System> {
        Box::new(self.clone())
    }
}

impl WritableSystem for InMemorySystem {
    fn create_new_file(&self, path: &SystemPath) -> Result<()> {
        self.memory_fs.create_new_file(path)
    }

    fn write_file_bytes(&self, path: &SystemPath, content: &[u8]) -> Result<()> {
        self.memory_fs.write_file(path, content)
    }

    fn create_directory_all(&self, path: &SystemPath) -> Result<()> {
        self.memory_fs.create_directory_all(path)
    }

    fn dyn_clone(&self) -> Box<dyn WritableSystem> {
        Box::new(self.clone())
    }
}