ephemeral_env/

lib.rs

//! A testing utility for creating ephemeral environments which are reset once they go out of
//! scope.
//!
//! ## Examples:
//!
//! Once an [EphemeralEnv] drops out of scope, new env vars which were created whilst it was
//! in scope will be dropped.
//!
//! ```
//! use ephemeral_env::EphemeralEnv;
//!
//! assert!(std::env::var("MY_ENVIRONMENT_VARIABLE").is_err());
//! #[cfg(feature = "sync")]
//! {
//!     let mut ephemeral = EphemeralEnv::from_env_sync().unwrap();
//!     ephemeral.set_var("MY_ENVIRONMENT_VARIABLE", "test");
//!     assert_eq!(std::env::var("MY_ENVIRONMENT_VARIABLE").unwrap(), "test");
//! }
//! assert!(std::env::var("MY_ENVIRONMENT_VARIABLE").is_err());
//! ```
//!
//! Similarly, once an EphemeralEnv drops out of scope, modified env vars will revert to
//! the value they held when the EphemeralEnv was created.
//!
//! ```
//! use ephemeral_env::EphemeralEnv;
//!
//! unsafe {
//!     std::env::set_var("MY_ENVIRONMENT_VARIABLE", "Spanners! Shh!");
//! }
//! {
//!     let mut ephemeral = EphemeralEnv::from_env_sync().unwrap();
//!     ephemeral.set_var("MY_ENVIRONMENT_VARIABLE", "test");
//!     assert_eq!(std::env::var("MY_ENVIRONMENT_VARIABLE").unwrap(), "test");
//! }
//! assert_eq!(std::env::var("MY_ENVIRONMENT_VARIABLE").unwrap(), "Spanners! Shh!");
//!```
//!
//! ## Async interface
//!
//! If you're testing across await points then it's inadvisable to use
//! [EphemeralEnv::from_env_sync()], as this will use a [std::sync::Mutex], which isn't [Send].
//!
//! Instead, `ephemeral_env` provides [from_env_async()], which will use Tokio's async-compatible
//! [Mutex](tokio::sync::Mutex) instead:
//!
//! ```
//! use ephemeral_env::EphemeralEnv;
//!
//! # #[cfg(feature = "async")]
//! # #[tokio::test]
//! # async fn test_async_ephemeral_env() {
//! assert!(std::env::var("MY_ENVIRONMENT_VARIABLE").is_err());
//! {
//!     let mut ephemeral = EphemeralEnv::from_env_async().unwrap();
//!     ephemeral.set_var("MY_ENVIRONMENT_VARIABLE", "test");
//!     // Without the async interface, this call would cause a warning to be shown, and there
//!     // would be a potential for the mutex held by the ephemeral environment to block the
//!     // event loop.
//!     tokio::time::sleep(Duration::from_millis(500)).await;
//!     assert_eq!(std::env::var("MY_ENVIRONMENT_VARIABLE").unwrap(), "test");
//! }
//! assert!(std::env::var("MY_ENVIRONMENT_VARIABLE").is_err());
//! # }
//! ```
#[cfg(feature = "async")]
use once_cell::sync::Lazy;
#[cfg(feature = "sync")]
use std::sync::MutexGuard;
use std::sync::PoisonError;
use std::{collections::HashMap, env};
#[cfg(feature = "async")]
use tokio::sync::MutexGuard as AsyncMutexGuard;

#[cfg(feature = "async")]
static ASYNC_ENV_LOCK: Lazy<tokio::sync::Mutex<()>> = Lazy::new(|| tokio::sync::Mutex::new(()));
#[cfg(feature = "sync")]
static SYNC_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

pub struct EphemeralEnv {
    initial_vars: HashMap<String, String>,
    #[cfg(feature = "sync")]
    _sync_guard: Option<MutexGuard<'static, ()>>,
    #[cfg(feature = "async")]
    _async_guard: Option<AsyncMutexGuard<'static, ()>>,
}

impl EphemeralEnv {
    /// Create a new synchronous EphemeralEnv.
    ///
    /// Takes a copy of the current environment variables, and prevents other EphemeralEnvs from
    /// being started by holding a [MutexGuard](std::sync::MutexGuard).
    ///
    /// **Do not use this in tests of async code**. [std::sync::MutexGuard] is not [Send], and so
    /// holding onto one across an await point can result in deadlocks.
    #[cfg(feature = "sync")]
    pub fn from_env_sync() -> Result<Self, PoisonError<MutexGuard<'static, ()>>> {
        let guard = SYNC_ENV_LOCK.lock()?;
        Ok(Self {
            initial_vars: HashMap::from_iter(env::vars()),
            _sync_guard: Some(guard),
            #[cfg(feature = "async")]
            _async_guard: None,
        })
    }

    /// Create a new asynchronous EphemeralEnv.
    ///
    /// Takes a copy of the current environment variables, and prevents other EphemeralEnvs from
    /// being started by holding a [MutexGuard](tokio::sync::MutexGuard).
    #[cfg(feature = "async")]
    pub async fn from_env_async() -> Result<Self, PoisonError<AsyncMutexGuard<'static, ()>>> {
        let guard = ASYNC_ENV_LOCK.lock().await;
        Ok(Self {
            initial_vars: HashMap::from_iter(env::vars()),
            #[cfg(feature = "async")]
            _async_guard: Some(guard),
            #[cfg(feature = "sync")]
            _sync_guard: None,
        })
    }

    /// Set one or more environment variables from a Vec of string tuples.
    pub fn with_vars(&mut self, vars: Vec<(String, String)>) {
        for (key, value) in vars.iter() {
            unsafe {
                std::env::set_var(key, value);
            }
        }
    }

    /// Set a single env var.
    ///
    /// This is a convenience function to spare the caller from adding unsafe {}
    /// blocks hither and yon.
    pub fn set_var(&mut self, key: impl Into<String>, value: impl Into<String>) {
        unsafe {
            std::env::set_var(key.into(), value.into());
        }
    }
}

impl Drop for EphemeralEnv {
    fn drop(&mut self) {
        for (key, value) in std::env::vars() {
            unsafe {
                match self.initial_vars.get(&key) {
                    Some(initial) => {
                        if initial != &value {
                            std::env::set_var(key, initial);
                        }
                    }
                    None => std::env::remove_var(key),
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use once_cell::sync::Lazy;
    use tokio::sync::Mutex;

    static TEST_LOCK: Lazy<tokio::sync::Mutex<()>> = Lazy::new(|| Mutex::new(()));

    #[cfg(feature = "async")]
    #[tokio::test]
    async fn test_async_env_locking_does_not_block_the_event_loop() {
        let _guard = TEST_LOCK.lock().await;
        let key = "ASYNC_TEST_KEY";
        unsafe {
            env::set_var(key, "original");
        }

        {
            let mut env = EphemeralEnv::from_env_async().await.unwrap();
            env.set_var(key, "temporary");

            // We can safely await other things while holding the lock
            // without breaking the executor or the compiler.
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

            assert_eq!(env::var(key).unwrap(), "temporary");
        } // Drops here, releasing lock

        assert_eq!(env::var(key).unwrap(), "original");
        unsafe {
            env::remove_var(key);
        }
    }

    #[cfg(feature = "sync")]
    #[tokio::test]
    async fn test_restores_modified_variable_on_drop() {
        let _guard = TEST_LOCK.lock().await;
        let key = "TEST_ENV_AUTO_LOCK";
        unsafe {
            std::env::set_var(key, "original");
        }

        {
            let _ephemeral = EphemeralEnv::from_env_sync().unwrap();
            unsafe {
                std::env::set_var(key, "temp");
            }
            assert_eq!(std::env::var(key).unwrap(), "temp");
        }

        assert_eq!(std::env::var(key).unwrap(), "original");
        unsafe {
            std::env::remove_var(key);
        }
    }

    #[cfg(feature = "sync")]
    #[tokio::test]
    async fn test_sync_ephemeral_env_is_threadsafe() {
        let _guard = TEST_LOCK.lock().await;
        // Spawn some threads that would normally race. If the lock works, they will run
        // sequentially without error.
        let handles: Vec<_> = (0..10)
            .map(|i| {
                std::thread::spawn(move || {
                    let key = "TEST_RACE_CONDITION";
                    let _ephemeral = EphemeralEnv::from_env_sync().unwrap();

                    // If race conditions were present, another thread might
                    // change this value while we are sleeping.
                    let val = format!("val_{}", i);
                    unsafe {
                        std::env::set_var(key, &val);
                    }

                    // Simulate some work
                    std::thread::sleep(std::time::Duration::from_millis(10));
                    assert_eq!(std::env::var(key).unwrap(), val);
                })
            })
            .collect();

        for handle in handles {
            handle.join().unwrap();
        }
    }

    #[cfg(feature = "async")]
    #[tokio::test]
    async fn test_async_ephemeral_env_is_threadsafe() {
        let _guard = TEST_LOCK.lock().await;
        // Spawn some threads that would normally race. If the lock works, they will run
        // sequentially without error.
        let handles: Vec<_> = (0..10)
            .map(|i| {
                tokio::spawn(async move {
                    let key = "TEST_RACE_CONDITION";
                    let _ephemeral = EphemeralEnv::from_env_async().await.unwrap();

                    // If race conditions were present, another thread might
                    // change this value while we are sleeping.
                    let val = format!("val_{}", i);
                    unsafe {
                        std::env::set_var(key, &val);
                    }

                    // Simulate some work
                    std::thread::sleep(std::time::Duration::from_millis(10));
                    assert_eq!(std::env::var(key).unwrap(), val);
                })
            })
            .collect();

        for handle in handles {
            handle.await.unwrap()
        }
    }
}