static TEST_ENV_MUTEX: std::sync::OnceLock<std::sync::Mutex<()>> = std::sync::OnceLock::new();
thread_local! {
static TEST_ENV_DEPTH: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
static TEST_ENV_HELD: std::cell::RefCell<Option<std::sync::MutexGuard<'static, ()>>> =
const { std::cell::RefCell::new(None) };
}
pub fn test_env_lock() -> TestEnvGuard {
let mutex = TEST_ENV_MUTEX.get_or_init(|| std::sync::Mutex::new(()));
TEST_ENV_DEPTH.with(|depth| {
if depth.get() == 0 {
let guard = mutex
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
TEST_ENV_HELD.with(|held| *held.borrow_mut() = Some(guard));
}
depth.set(depth.get() + 1);
});
TestEnvGuard {
_not_send_sync: std::marker::PhantomData,
}
}
#[cfg(any(test, feature = "test-support"))]
pub fn holds_test_env_lock() -> bool {
TEST_ENV_DEPTH.with(|depth| depth.get() > 0)
}
pub struct TestEnvGuard {
_not_send_sync: std::marker::PhantomData<*const ()>,
}
impl Drop for TestEnvGuard {
fn drop(&mut self) {
TEST_ENV_DEPTH.with(|depth| {
let next = depth.get().saturating_sub(1);
depth.set(next);
if next == 0 {
TEST_ENV_HELD.with(|held| *held.borrow_mut() = None);
}
});
}
}
#[cfg(any(test, feature = "test-support"))]
pub fn set_var<K: AsRef<std::ffi::OsStr>, V: AsRef<std::ffi::OsStr>>(key: K, value: V) {
assert_lock_held();
unsafe { std::env::set_var(key, value) };
}
#[cfg(any(test, feature = "test-support"))]
pub fn remove_var<K: AsRef<std::ffi::OsStr>>(key: K) {
assert_lock_held();
unsafe { std::env::remove_var(key) };
}
#[cfg(any(test, feature = "test-support"))]
fn assert_lock_held() {
assert!(
holds_test_env_lock(),
"test environment mutation requires the current thread to hold test_env_lock()"
);
}