use std::path::Path;
pub struct EnvGuard {
key: &'static str,
prev: Option<String>,
}
impl EnvGuard {
pub(crate) fn set(key: &'static str, val: impl AsRef<Path>) -> Self {
let prev = std::env::var(key).ok();
unsafe { std::env::set_var(key, val.as_ref().as_os_str()) };
Self { key, prev }
}
pub(crate) fn remove(key: &'static str) -> Self {
let prev = std::env::var(key).ok();
unsafe { std::env::remove_var(key) };
Self { key, prev }
}
}
impl Drop for EnvGuard {
fn drop(&mut self) {
match &self.prev {
Some(v) => unsafe { std::env::set_var(self.key, v) },
None => unsafe { std::env::remove_var(self.key) },
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serial_test::serial;
#[test]
#[serial]
fn envguard_set_and_restore_when_unset() {
let key = "__TEST_AGENTIC_CONFIG_ENVVAR_A";
let _r = EnvGuard::remove(key);
{
let _g = EnvGuard::set(key, "/tmp/test");
assert_eq!(std::env::var(key).unwrap(), "/tmp/test");
}
assert!(std::env::var(key).is_err(), "should restore to unset");
}
#[test]
#[serial]
fn envguard_restore_previous_value() {
let key = "__TEST_AGENTIC_CONFIG_ENVVAR_B";
let _orig = EnvGuard::set(key, "/original");
{
let _g = EnvGuard::set(key, "/shadow");
assert_eq!(std::env::var(key).unwrap(), "/shadow");
}
assert_eq!(std::env::var(key).unwrap(), "/original");
}
#[test]
#[serial]
fn envguard_remove_and_restore() {
let key = "__TEST_AGENTIC_CONFIG_ENVVAR_C";
let _orig = EnvGuard::set(key, "/value");
{
let _g = EnvGuard::remove(key);
assert!(std::env::var(key).is_err());
}
assert_eq!(std::env::var(key).unwrap(), "/value");
}
}