use std::collections::BTreeMap;
use std::ffi::{OsStr, OsString};
use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
use std::path::PathBuf;
use std::sync::{Mutex, MutexGuard};
static CWD_MUTEX: Mutex<()> = Mutex::new(());
static ENV_MUTEX: Mutex<()> = Mutex::new(());
pub fn lock_env() -> EnvGuard {
let lock = ENV_MUTEX.lock().unwrap_or_else(|e| e.into_inner());
EnvGuard {
_lock: lock,
originals: BTreeMap::new(),
}
}
pub fn lock_env_and_cwd() -> EnvCwdGuard {
let env = lock_env();
let cwd = lock_cwd();
EnvCwdGuard {
env: Some(env),
cwd: Some(cwd),
}
}
pub fn lock_cwd() -> CwdGuard {
let lock = CWD_MUTEX.lock().unwrap_or_else(|e| e.into_inner());
let original = std::env::current_dir().expect("failed to get current dir");
CwdGuard {
_lock: lock,
original,
}
}
pub struct CwdGuard {
_lock: MutexGuard<'static, ()>,
original: PathBuf,
}
pub struct EnvCwdGuard {
env: Option<EnvGuard>,
cwd: Option<CwdGuard>,
}
impl EnvCwdGuard {
pub fn env(&mut self) -> &mut EnvGuard {
self.env
.as_mut()
.expect("combined test guard environment guard was already dropped")
}
}
impl Drop for EnvCwdGuard {
fn drop(&mut self) {
let cwd_panic = self
.cwd
.take()
.and_then(|cwd| catch_unwind(AssertUnwindSafe(|| drop(cwd))).err());
if let Some(env) = self.env.take() {
drop(env);
}
if let Some(payload) = cwd_panic {
resume_unwind(payload);
}
}
}
impl Drop for CwdGuard {
fn drop(&mut self) {
if let Err(err) = std::env::set_current_dir(&self.original) {
if std::thread::panicking() {
eprintln!(
"CwdGuard: failed to restore original working directory {:?}: {}",
self.original, err
);
} else {
panic!(
"CwdGuard: failed to restore original working directory {:?}: {}",
self.original, err
);
}
}
}
}
pub struct EnvGuard {
_lock: MutexGuard<'static, ()>,
originals: BTreeMap<OsString, Option<OsString>>,
}
impl EnvGuard {
pub fn set(&mut self, name: impl AsRef<OsStr>, value: impl AsRef<OsStr>) {
let name = name.as_ref().to_os_string();
self.originals
.entry(name.clone())
.or_insert_with(|| std::env::var_os(&name));
unsafe { std::env::set_var(&name, value) };
}
pub fn unset(&mut self, name: impl AsRef<OsStr>) {
let name = name.as_ref().to_os_string();
self.originals
.entry(name.clone())
.or_insert_with(|| std::env::var_os(&name));
unsafe { std::env::remove_var(&name) };
}
}
impl Drop for EnvGuard {
fn drop(&mut self) {
unsafe {
for (name, original) in &self.originals {
if let Some(value) = original {
std::env::set_var(name, value);
} else {
std::env::remove_var(name);
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn lock_env_restores_multiple_variables() {
let first = "ARF_TEST_UTILS_MULTIPLE_FIRST";
let second = "ARF_TEST_UTILS_MULTIPLE_SECOND";
let (original_first, original_second) = {
let _guard = lock_env();
(std::env::var_os(first), std::env::var_os(second))
};
{
let mut guard = lock_env();
guard.set(first, "changed-first");
guard.set(second, "changed-second");
assert_eq!(
std::env::var_os(first),
Some(OsString::from("changed-first"))
);
assert_eq!(
std::env::var_os(second),
Some(OsString::from("changed-second"))
);
}
let _guard = lock_env();
assert_eq!(std::env::var_os(first), original_first);
assert_eq!(std::env::var_os(second), original_second);
}
#[test]
fn lock_env_restores_first_original_value_after_repeated_mutations() {
let name = "ARF_TEST_UTILS_REPEATED";
let original = {
let _guard = lock_env();
std::env::var_os(name)
};
{
let mut guard = lock_env();
guard.set(name, "first-change");
guard.unset(name);
guard.set(name, "second-change");
}
let _guard = lock_env();
assert_eq!(std::env::var_os(name), original);
}
#[test]
fn lock_env_supports_mixed_set_and_unset_mutations() {
let set_name = "ARF_TEST_UTILS_MIXED_SET";
let unset_name = "ARF_TEST_UTILS_MIXED_UNSET";
let (original_set, original_unset) = {
let _guard = lock_env();
(std::env::var_os(set_name), std::env::var_os(unset_name))
};
{
let mut guard = lock_env();
guard.set(set_name, "set-value");
guard.unset(unset_name);
assert_eq!(
std::env::var_os(set_name),
Some(OsString::from("set-value"))
);
assert_eq!(std::env::var_os(unset_name), None);
}
let _guard = lock_env();
assert_eq!(std::env::var_os(set_name), original_set);
assert_eq!(std::env::var_os(unset_name), original_unset);
}
#[test]
fn lock_env_restores_a_variable_that_was_originally_unset() {
let name = format!("ARF_TEST_UTILS_ORIGINALLY_UNSET_{}", std::process::id());
let name = name.as_str();
{
let mut guard = lock_env();
assert_eq!(std::env::var_os(name), None);
guard.set(name, "temporary-value");
}
let _guard = lock_env();
assert_eq!(std::env::var_os(name), None);
}
#[test]
fn lock_env_allows_mutating_two_variables_without_deadlocking() {
let first = "ARF_TEST_UTILS_NO_DEADLOCK_FIRST";
let second = "ARF_TEST_UTILS_NO_DEADLOCK_SECOND";
let (original_first, original_second) = {
let _guard = lock_env();
(std::env::var_os(first), std::env::var_os(second))
};
{
let mut guard = lock_env();
guard.set(first, "first-value");
guard.set(second, "second-value");
}
let _guard = lock_env();
assert_eq!(std::env::var_os(first), original_first);
assert_eq!(std::env::var_os(second), original_second);
}
#[test]
fn lock_env_and_cwd_restores_both_after_combined_use() {
let name = "ARF_TEST_UTILS_COMBINED";
let temp_dir = tempfile::tempdir().unwrap();
let (original_value, original_cwd) = {
let _guard = lock_env_and_cwd();
(std::env::var_os(name), std::env::current_dir().unwrap())
};
{
let mut guard = lock_env_and_cwd();
guard.env().set(name, "combined-value");
std::env::set_current_dir(temp_dir.path()).unwrap();
assert_eq!(
std::env::var_os(name),
Some(OsString::from("combined-value"))
);
assert_eq!(
std::env::current_dir().unwrap().canonicalize().ok(),
temp_dir.path().canonicalize().ok()
);
}
let _guard = lock_env_and_cwd();
assert_eq!(std::env::var_os(name), original_value);
assert_eq!(std::env::current_dir().unwrap(), original_cwd);
}
}