51 lines
1.2 KiB
Rust
51 lines
1.2 KiB
Rust
use std::collections::BTreeMap;
|
|
use std::env;
|
|
use std::ffi::{OsStr, OsString};
|
|
use std::sync::{Mutex, MutexGuard, OnceLock};
|
|
|
|
fn env_lock() -> &'static Mutex<()> {
|
|
static ENV_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
|
ENV_LOCK.get_or_init(|| Mutex::new(()))
|
|
}
|
|
|
|
pub(crate) struct TestEnv {
|
|
_guard: MutexGuard<'static, ()>,
|
|
saved: BTreeMap<String, Option<OsString>>,
|
|
}
|
|
|
|
impl TestEnv {
|
|
pub(crate) fn new() -> Self {
|
|
Self {
|
|
_guard: env_lock().lock().expect("test env lock poisoned"),
|
|
saved: BTreeMap::new(),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn set_var(&mut self, key: &str, value: impl AsRef<OsStr>) {
|
|
self.snapshot(key);
|
|
unsafe {
|
|
env::set_var(key, value);
|
|
}
|
|
}
|
|
|
|
fn snapshot(&mut self, key: &str) {
|
|
self.saved
|
|
.entry(key.to_string())
|
|
.or_insert_with(|| env::var_os(key));
|
|
}
|
|
}
|
|
|
|
impl Drop for TestEnv {
|
|
fn drop(&mut self) {
|
|
for (key, value) in &self.saved {
|
|
unsafe {
|
|
if let Some(value) = value {
|
|
env::set_var(key, value);
|
|
} else {
|
|
env::remove_var(key);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|