diff --git a/src/commands/deploy.rs b/src/commands/deploy.rs index ff3878d..ba3a589 100644 --- a/src/commands/deploy.rs +++ b/src/commands/deploy.rs @@ -117,6 +117,7 @@ pub async fn rollback_latest() -> Result> { #[cfg(test)] mod tests { use super::*; + use crate::test_utils::EnvGuard; use std::sync::{Mutex, OnceLock}; fn env_lock() -> &'static Mutex<()> { @@ -172,18 +173,18 @@ mod tests { #[tokio::test] async fn load_manifest_nonexistent_returns_default() { - let _guard = env_lock().lock().expect("env lock poisoned"); - std::env::set_var("UPDATE_STORAGE_PATH", "/tmp/status-test-nonexistent-path"); + let _lock = env_lock().lock().expect("env lock poisoned"); + let dir = tempfile::tempdir().unwrap(); + let _env = EnvGuard::set("UPDATE_STORAGE_PATH", dir.path().to_str().unwrap()); let manifest = load_manifest().await.unwrap(); assert!(manifest.entries.is_empty()); - std::env::remove_var("UPDATE_STORAGE_PATH"); } #[tokio::test] async fn save_and_load_manifest_roundtrip() { - let _guard = env_lock().lock().expect("env lock poisoned"); + let _lock = env_lock().lock().expect("env lock poisoned"); let dir = tempfile::tempdir().unwrap(); - std::env::set_var("UPDATE_STORAGE_PATH", dir.path().to_str().unwrap()); + let _env = EnvGuard::set("UPDATE_STORAGE_PATH", dir.path().to_str().unwrap()); let manifest = RollbackManifest { entries: vec![RollbackEntry { @@ -198,15 +199,13 @@ mod tests { let loaded = load_manifest().await.unwrap(); assert_eq!(loaded.entries.len(), 1); assert_eq!(loaded.entries[0].job_id, "test-job"); - - std::env::remove_var("UPDATE_STORAGE_PATH"); } #[tokio::test] async fn record_rollback_appends_entry() { - let _guard = env_lock().lock().expect("env lock poisoned"); + let _lock = env_lock().lock().expect("env lock poisoned"); let dir = tempfile::tempdir().unwrap(); - std::env::set_var("UPDATE_STORAGE_PATH", dir.path().to_str().unwrap()); + let _env = EnvGuard::set("UPDATE_STORAGE_PATH", dir.path().to_str().unwrap()); // Save an initial empty manifest save_manifest(&RollbackManifest::default()).await.unwrap(); @@ -222,15 +221,13 @@ mod tests { assert_eq!(loaded.entries.len(), 2); assert_eq!(loaded.entries[0].job_id, "job-1"); assert_eq!(loaded.entries[1].job_id, "job-2"); - - std::env::remove_var("UPDATE_STORAGE_PATH"); } #[tokio::test] async fn backup_current_binary_creates_file() { - let _guard = env_lock().lock().expect("env lock poisoned"); + let _lock = env_lock().lock().expect("env lock poisoned"); let dir = tempfile::tempdir().unwrap(); - std::env::set_var("UPDATE_STORAGE_PATH", dir.path().to_str().unwrap()); + let _env = EnvGuard::set("UPDATE_STORAGE_PATH", dir.path().to_str().unwrap()); // Create a fake binary to back up let src = dir.path().join("status"); @@ -243,20 +240,16 @@ mod tests { let content = tokio::fs::read(&backup_path).await.unwrap(); assert_eq!(content, b"fake binary content"); - - std::env::remove_var("UPDATE_STORAGE_PATH"); } #[tokio::test] async fn rollback_latest_with_empty_manifest_returns_none() { - let _guard = env_lock().lock().expect("env lock poisoned"); + let _lock = env_lock().lock().expect("env lock poisoned"); let dir = tempfile::tempdir().unwrap(); - std::env::set_var("UPDATE_STORAGE_PATH", dir.path().to_str().unwrap()); + let _env = EnvGuard::set("UPDATE_STORAGE_PATH", dir.path().to_str().unwrap()); save_manifest(&RollbackManifest::default()).await.unwrap(); let result = rollback_latest().await.unwrap(); assert!(result.is_none()); - - std::env::remove_var("UPDATE_STORAGE_PATH"); } } diff --git a/src/commands/version_check.rs b/src/commands/version_check.rs index 8b85620..b12facd 100644 --- a/src/commands/version_check.rs +++ b/src/commands/version_check.rs @@ -38,6 +38,7 @@ pub async fn check_remote_version() -> Result> { #[cfg(test)] mod tests { use super::*; + use crate::test_utils::EnvGuard; use std::sync::{Mutex, OnceLock}; fn env_lock() -> &'static Mutex<()> { @@ -78,18 +79,17 @@ mod tests { #[tokio::test] async fn check_remote_version_no_env_returns_none() { - let _guard = env_lock().lock().expect("env lock poisoned"); - std::env::remove_var("UPDATE_SERVER_URL"); + let _lock = env_lock().lock().expect("env lock poisoned"); + let _env = EnvGuard::remove("UPDATE_SERVER_URL"); let result = check_remote_version().await.unwrap(); assert!(result.is_none()); } #[tokio::test] async fn check_remote_version_empty_env_returns_none() { - let _guard = env_lock().lock().expect("env lock poisoned"); - std::env::set_var("UPDATE_SERVER_URL", ""); + let _lock = env_lock().lock().expect("env lock poisoned"); + let _env = EnvGuard::set("UPDATE_SERVER_URL", ""); let result = check_remote_version().await.unwrap(); assert!(result.is_none()); - std::env::remove_var("UPDATE_SERVER_URL"); } } diff --git a/src/lib.rs b/src/lib.rs index ace8c72..2451389 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,5 +7,8 @@ pub mod security; pub mod transport; pub mod utils; +#[cfg(test)] +pub(crate) mod test_utils; + // Crate version exposed for runtime queries pub const VERSION: &str = env!("CARGO_PKG_VERSION"); diff --git a/src/security/scopes.rs b/src/security/scopes.rs index 947277a..8b38b0d 100644 --- a/src/security/scopes.rs +++ b/src/security/scopes.rs @@ -32,6 +32,7 @@ impl Scopes { #[cfg(test)] mod tests { use super::*; + use crate::test_utils::EnvGuard; use std::sync::{Mutex, OnceLock}; fn env_lock() -> &'static Mutex<()> { @@ -62,52 +63,48 @@ mod tests { #[test] fn scopes_from_env_parses_comma_separated() { - let _guard = env_lock().lock().expect("env lock poisoned"); - std::env::set_var("AGENT_SCOPES", "docker:restart,docker:logs,admin"); + let _lock = env_lock().lock().expect("env lock poisoned"); + let _env = EnvGuard::set("AGENT_SCOPES", "docker:restart,docker:logs,admin"); let scopes = Scopes::from_env(); assert!(scopes.is_allowed("docker:restart")); assert!(scopes.is_allowed("docker:logs")); assert!(scopes.is_allowed("admin")); assert!(!scopes.is_allowed("docker:stop")); - std::env::remove_var("AGENT_SCOPES"); } #[test] fn scopes_from_env_trims_whitespace() { - let _guard = env_lock().lock().expect("env lock poisoned"); - std::env::set_var("AGENT_SCOPES", " docker:restart , admin "); + let _lock = env_lock().lock().expect("env lock poisoned"); + let _env = EnvGuard::set("AGENT_SCOPES", " docker:restart , admin "); let scopes = Scopes::from_env(); assert!(scopes.is_allowed("docker:restart")); assert!(scopes.is_allowed("admin")); - std::env::remove_var("AGENT_SCOPES"); } #[test] fn scopes_from_env_skips_empty_items() { - let _guard = env_lock().lock().expect("env lock poisoned"); - std::env::set_var("AGENT_SCOPES", "docker:restart,,, ,admin"); + let _lock = env_lock().lock().expect("env lock poisoned"); + let _env = EnvGuard::set("AGENT_SCOPES", "docker:restart,,, ,admin"); let scopes = Scopes::from_env(); assert!(scopes.is_allowed("docker:restart")); assert!(scopes.is_allowed("admin")); // The empty strings should NOT be in the set assert!(!scopes.is_allowed("")); - std::env::remove_var("AGENT_SCOPES"); } #[test] fn scopes_from_env_missing_var_allows_all() { - let _guard = env_lock().lock().expect("env lock poisoned"); - std::env::remove_var("AGENT_SCOPES"); + let _lock = env_lock().lock().expect("env lock poisoned"); + let _env = EnvGuard::remove("AGENT_SCOPES"); let scopes = Scopes::from_env(); assert!(scopes.is_allowed("anything")); } #[test] fn scopes_from_env_empty_string_allows_all() { - let _guard = env_lock().lock().expect("env lock poisoned"); - std::env::set_var("AGENT_SCOPES", ""); + let _lock = env_lock().lock().expect("env lock poisoned"); + let _env = EnvGuard::set("AGENT_SCOPES", ""); let scopes = Scopes::from_env(); assert!(scopes.is_allowed("anything")); - std::env::remove_var("AGENT_SCOPES"); } } diff --git a/src/test_utils.rs b/src/test_utils.rs new file mode 100644 index 0000000..f3c7c30 --- /dev/null +++ b/src/test_utils.rs @@ -0,0 +1,31 @@ +/// A drop-based guard that saves an environment variable's original value before modification +/// and restores it when dropped, ensuring cleanup even when test assertions panic. +pub(crate) struct EnvGuard { + key: &'static str, + original: Option, +} + +impl EnvGuard { + /// Sets `key` to `value` and saves the previous value for restoration on drop. + pub(crate) fn set(key: &'static str, value: &str) -> Self { + let original = std::env::var(key).ok(); + std::env::set_var(key, value); + Self { key, original } + } + + /// Removes `key` from the environment and saves the previous value for restoration on drop. + pub(crate) fn remove(key: &'static str) -> Self { + let original = std::env::var(key).ok(); + std::env::remove_var(key); + Self { key, original } + } +} + +impl Drop for EnvGuard { + fn drop(&mut self) { + match &self.original { + Some(v) => std::env::set_var(self.key, v), + None => std::env::remove_var(self.key), + } + } +}