refactor: Add EnvGuard RAII wrapper for safe env var handling in tests (#179)

ARCHITECTURE.md

@@ -744,6 +744,46 @@ User Input → CLI Parser → Mode Detection → Node Resolution
 - **Timeout handling**: Configurable per-connection and per-command
 - **Signal handling**: Clean shutdown on Ctrl+C with two-stage confirmation
 
+### Test Environment-Variable Mutation Pattern (`EnvGuard`)
+
+Several test suites must temporarily set or remove process-wide environment
+variables (e.g. `BACKENDAI_CLUSTER_HOSTS`, `HOME`, `SSH_AUTH_SOCK`). Under
+Rust 2024 edition, `std::env::set_var` and `std::env::remove_var` are marked
+`unsafe` because concurrent mutation of the environment is undefined behaviour
+at the libc level on glibc, musl, and macOS. `EnvGuard` centralises all such
+mutations in `src/test_helpers/env_guard.rs`.
+
+**Soundness contract**: every test that constructs an `EnvGuard` MUST be
+annotated with `#[serial_test::serial]`. Every other test in the same crate
+binary that reads or mutates the same variable MUST also carry `#[serial]` (or
+a matching `#[serial(key)]` group). Note that `#[serial]` only serializes
+against other `#[serial]` / `#[parallel]` tests — unannotated tests may still
+run concurrently with serial ones and would race on environment reads. This is
+not an `EnvGuard` limitation; it is an inherent constraint of the libc
+environment-variable API.
+
+```rust
+use serial_test::serial;
+use crate::test_helpers::EnvGuard;
+
+#[test]
+#[serial]
+fn my_test() {
+    let _host = EnvGuard::set("BACKENDAI_CLUSTER_HOSTS", "node1,node2");
+    // Variable is automatically restored when `_host` drops at end of scope.
+}
+```
+
+**Integration tests** access the same struct via a `#[path]`-based re-export
+in `tests/common/mod.rs`, which avoids making `EnvGuard` part of the public
+`bssh` crate API while keeping a single source of truth. When adding a new
+integration-test binary that needs `EnvGuard`, add `mod common;` at the top of
+that file and use `common::EnvGuard`.
+
+Use `#[serial(key)]` (a named group) when two sets of tests touch different,
+non-overlapping variables and can therefore run concurrently with each other
+but not with themselves; omit the key (plain `#[serial]`) when in doubt.
+
 ## Security Model
 
 ### Authentication

src/cli/mode_detection_tests.rs

@@ -20,47 +20,32 @@
 #[cfg(test)]
 mod tests {
     use crate::cli::pdsh::PDSH_COMPAT_ENV_VAR;
+    use crate::test_helpers::EnvGuard;
     use serial_test::serial;
     use std::env;
 
     /// Test that environment variable detection works for "1"
     #[test]
     #[serial]
     fn test_env_var_detection_one() {
-        // Save and restore env var state
-        let original = env::var(PDSH_COMPAT_ENV_VAR).ok();
-
-        env::set_var(PDSH_COMPAT_ENV_VAR, "1");
+        let _guard = EnvGuard::set(PDSH_COMPAT_ENV_VAR, "1");
 
         // Create a test for the env var checking logic
         let value = env::var(PDSH_COMPAT_ENV_VAR).ok();
         assert!(value.is_some());
         let value = value.unwrap();
         assert!(value == "1" || value.to_lowercase() == "true");
-
-        // Restore
-        match original {
-            Some(v) => env::set_var(PDSH_COMPAT_ENV_VAR, v),
-            None => env::remove_var(PDSH_COMPAT_ENV_VAR),
-        }
     }
 
     /// Test that environment variable detection works for "true"
     #[test]
     #[serial]
     fn test_env_var_detection_true() {
-        let original = env::var(PDSH_COMPAT_ENV_VAR).ok();
-
-        env::set_var(PDSH_COMPAT_ENV_VAR, "true");
+        let _guard = EnvGuard::set(PDSH_COMPAT_ENV_VAR, "true");
 
         let value = env::var(PDSH_COMPAT_ENV_VAR).ok();
         assert!(value.is_some());
         assert_eq!(value.unwrap().to_lowercase(), "true");
-
-        match original {
-            Some(v) => env::set_var(PDSH_COMPAT_ENV_VAR, v),
-            None => env::remove_var(PDSH_COMPAT_ENV_VAR),
-        }
     }
 
     /// Test that environment variable detection works for "TRUE" (case insensitive)
@@ -84,47 +69,38 @@ mod tests {
     #[test]
     #[serial]
     fn test_env_var_not_set() {
-        let original = env::var(PDSH_COMPAT_ENV_VAR).ok();
-
-        env::remove_var(PDSH_COMPAT_ENV_VAR);
+        let _guard = EnvGuard::remove(PDSH_COMPAT_ENV_VAR);
 
         let value = env::var(PDSH_COMPAT_ENV_VAR).ok();
         assert!(value.is_none());
-
-        // Restore
-        if let Some(v) = original {
-            env::set_var(PDSH_COMPAT_ENV_VAR, v);
-        }
     }
 
     /// Test that invalid env var values are not treated as enabled
     #[test]
     #[serial]
     fn test_env_var_invalid_values() {
-        let original = env::var(PDSH_COMPAT_ENV_VAR).ok();
-
         // Test "0"
-        env::set_var(PDSH_COMPAT_ENV_VAR, "0");
-        let value = env::var(PDSH_COMPAT_ENV_VAR).unwrap();
-        let enabled = value == "1" || value.to_lowercase() == "true";
-        assert!(!enabled);
+        {
+            let _guard = EnvGuard::set(PDSH_COMPAT_ENV_VAR, "0");
+            let value = env::var(PDSH_COMPAT_ENV_VAR).unwrap();
+            let enabled = value == "1" || value.to_lowercase() == "true";
+            assert!(!enabled);
+        }
 
         // Test "false"
-        env::set_var(PDSH_COMPAT_ENV_VAR, "false");
-        let value = env::var(PDSH_COMPAT_ENV_VAR).unwrap();
-        let enabled = value == "1" || value.to_lowercase() == "true";
-        assert!(!enabled);
+        {
+            let _guard = EnvGuard::set(PDSH_COMPAT_ENV_VAR, "false");
+            let value = env::var(PDSH_COMPAT_ENV_VAR).unwrap();
+            let enabled = value == "1" || value.to_lowercase() == "true";
+            assert!(!enabled);
+        }
 
         // Test empty string
-        env::set_var(PDSH_COMPAT_ENV_VAR, "");
-        let value = env::var(PDSH_COMPAT_ENV_VAR).unwrap();
-        let enabled = value == "1" || value.to_lowercase() == "true";
-        assert!(!enabled);
-
-        // Restore
-        match original {
-            Some(v) => env::set_var(PDSH_COMPAT_ENV_VAR, v),
-            None => env::remove_var(PDSH_COMPAT_ENV_VAR),
+        {
+            let _guard = EnvGuard::set(PDSH_COMPAT_ENV_VAR, "");
+            let value = env::var(PDSH_COMPAT_ENV_VAR).unwrap();
+            let enabled = value == "1" || value.to_lowercase() == "true";
+            assert!(!enabled);
         }
     }
 

src/cli/pdsh.rs

@@ -141,12 +141,11 @@ pub const PDSH_COMPAT_ENV_VAR: &str = "BSSH_PDSH_COMPAT";
 ///
 /// # Examples
 ///
-/// ```
-/// use std::env;
-///
-/// // When environment variable is set
-/// env::set_var("BSSH_PDSH_COMPAT", "1");
-/// // is_pdsh_compat_mode() would return true
+/// ```ignore
+/// // When BSSH_PDSH_COMPAT=1 is set in the environment,
+/// // is_pdsh_compat_mode() returns true.
+/// // (Example is marked `ignore` because mutating environment variables in
+/// // Rust 2024 requires `unsafe`; see `EnvGuard` in `src/test_helpers`.)
 /// ```
 pub fn is_pdsh_compat_mode() -> bool {
     // Check environment variable first

src/config/tests.rs

@@ -18,13 +18,14 @@ use std::path::{Path, PathBuf};
 
 use super::types::{Config, InteractiveMode, NodeConfig};
 use super::utils::{expand_env_vars, expand_tilde};
+use crate::test_helpers::EnvGuard;
+use serial_test::serial;
 
 #[test]
+#[serial]
 fn test_expand_env_vars() {
-    unsafe {
-        std::env::set_var("TEST_VAR", "test_value");
-        std::env::set_var("TEST_USER", "testuser");
-    }
+    let _test_var = EnvGuard::set("TEST_VAR", "test_value");
+    let _test_user = EnvGuard::set("TEST_USER", "testuser");
 
     // Test ${VAR} syntax
     assert_eq!(expand_env_vars("Hello ${TEST_VAR}!"), "Hello test_value!");
@@ -49,23 +50,13 @@ fn test_expand_env_vars() {
 }
 
 #[test]
+#[serial]
 fn test_expand_tilde() {
-    // Save original HOME value
-    let original_home = std::env::var("HOME").ok();
-
-    // Set test HOME value
-    std::env::set_var("HOME", "/home/user");
+    let _home = EnvGuard::set("HOME", "/home/user");
 
     let path = Path::new("~/.ssh/config");
     let expanded = expand_tilde(path);
 
-    // Restore original HOME value
-    if let Some(home) = original_home {
-        std::env::set_var("HOME", home);
-    } else {
-        std::env::remove_var("HOME");
-    }
-
     assert_eq!(expanded, PathBuf::from("/home/user/.ssh/config"));
 }
 
@@ -194,14 +185,13 @@ clusters:
 }
 
 #[test]
+#[serial]
 fn test_backendai_env_parsing() {
-    // Set up Backend.AI environment variables
-    unsafe {
-        std::env::set_var("BACKENDAI_CLUSTER_HOSTS", "sub1,main1");
-        std::env::set_var("BACKENDAI_CLUSTER_HOST", "main1");
-        std::env::set_var("BACKENDAI_CLUSTER_ROLE", "main");
-        std::env::set_var("USER", "testuser");
-    }
+    // Set up Backend.AI environment variables; guards restore prior values on drop.
+    let _hosts = EnvGuard::set("BACKENDAI_CLUSTER_HOSTS", "sub1,main1");
+    let _host = EnvGuard::set("BACKENDAI_CLUSTER_HOST", "main1");
+    let _role = EnvGuard::set("BACKENDAI_CLUSTER_ROLE", "main");
+    let _user = EnvGuard::set("USER", "testuser");
 
     let cluster = Config::from_backendai_env().unwrap();
 
@@ -216,10 +206,11 @@ fn test_backendai_env_parsing() {
         _ => panic!("Expected Simple node config"),
     }
 
-    // Test with sub role - should skip the first (main) node
-    unsafe {
-        std::env::set_var("BACKENDAI_CLUSTER_ROLE", "sub");
-    }
+    // Test with sub role - should skip the first (main) node.
+    // The new guard shadows the outer `_role` binding; when it drops at the end
+    // of this test it'll restore to "main" (the value saved by `_role`), and
+    // then the outer `_role` drops and restores to the pre-test value.
+    let _role_sub = EnvGuard::set("BACKENDAI_CLUSTER_ROLE", "sub");
     let cluster = Config::from_backendai_env().unwrap();
     assert_eq!(cluster.nodes.len(), 1);
 
@@ -229,13 +220,6 @@ fn test_backendai_env_parsing() {
         }
         _ => panic!("Expected Simple node config"),
     }
-
-    // Clean up
-    unsafe {
-        std::env::remove_var("BACKENDAI_CLUSTER_HOSTS");
-        std::env::remove_var("BACKENDAI_CLUSTER_HOST");
-        std::env::remove_var("BACKENDAI_CLUSTER_ROLE");
-    }
 }
 
 #[test]
@@ -465,12 +449,11 @@ clusters:
 }
 
 #[test]
+#[serial]
 fn test_jump_host_env_var_expansion() {
-    // Set up test environment variables
-    unsafe {
-        std::env::set_var("TEST_BASTION_HOST", "bastion.example.com");
-        std::env::set_var("TEST_BASTION_PORT", "2222");
-    }
+    // Set up test environment variables; guards restore prior values on drop.
+    let _host = EnvGuard::set("TEST_BASTION_HOST", "bastion.example.com");
+    let _port = EnvGuard::set("TEST_BASTION_PORT", "2222");
 
     let yaml = r#"
 defaults:
@@ -504,12 +487,6 @@ clusters:
         config.get_cluster_jump_host(Some("staging")),
         Some("bastion.example.com".to_string())
     );
-
-    // Clean up
-    unsafe {
-        std::env::remove_var("TEST_BASTION_HOST");
-        std::env::remove_var("TEST_BASTION_PORT");
-    }
 }
 
 #[test]