Simplify env overriding

breezy-team · Jul 29, 2024 · 549794c · 549794c
1 parent 108122a
commit 549794c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 43 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -49,6 +49,7 @@ pub mod revisionid;
 pub mod status;
 pub mod subversion;
 pub mod tags;
+#[cfg(test)]
 pub mod tests;
 pub mod transform;
 pub mod transport;

diff --git a/src/tests.rs b/src/tests.rs
@@ -1,4 +1,5 @@
 use pyo3::prelude::*;
+use std::collections::HashMap;
 use std::fs;
 use std::path::PathBuf;
 use tempfile::TempDir;
@@ -8,9 +9,7 @@ pub struct TestEnv {
     pub working_dir: PathBuf,
     pub home_dir: PathBuf,
     pub old_cwd: PathBuf,
-    pub old_home: Option<String>,
-    pub old_email: Option<String>,
-    pub old_brz_home: Option<String>,
+    pub old_env: HashMap<String, Option<String>>,
 }
 
 impl TestEnv {
@@ -20,10 +19,11 @@ impl TestEnv {
         fs::create_dir(&working_dir).unwrap();
         let home_dir = temp_dir.path().join("home");
         fs::create_dir(&home_dir).unwrap();
+        let mut old_env = HashMap::new();
         let old_cwd = std::env::current_dir().unwrap();
-        let old_home = std::env::var("HOME").ok();
-        let old_email = std::env::var("BRZ_EMAIL").ok();
-        let old_brz_home = std::env::var("BRZ_HOME").ok();
+        old_env.insert("HOME".to_string(), std::env::var("HOME").ok());
+        old_env.insert("BRZ_EMAIL".to_string(), std::env::var("BRZ_EMAIL").ok());
+        old_env.insert("BRZ_HOME".to_string(), std::env::var("BRZ_HOME").ok());
         let brz_email = "Joe Tester <joe@example.com>";
         let breezy_home = home_dir.join(".config/breezy");
         std::env::set_current_dir(&working_dir).unwrap();
@@ -54,51 +54,29 @@ email = Joe Tester <joe@example.com>
             home_dir,
             working_dir,
             old_cwd,
-            old_home,
-            old_email,
-            old_brz_home,
+            old_env,
         }
     }
 }
 
 impl Drop for TestEnv {
     fn drop(&mut self) {
-        if let Some(dir) = self.old_home.as_ref() {
-            std::env::set_var("HOME", dir);
-        } else {
-            std::env::remove_var("HOME");
-        }
-        if let Some(email) = self.old_email.as_ref() {
-            std::env::set_var("BRZ_EMAIL", email);
-        } else {
-            std::env::remove_var("BRZ_EMAIL");
-        }
-        if let Some(dir) = self.old_brz_home.as_ref() {
-            std::env::set_var("BRZ_HOME", dir);
-        } else {
-            std::env::remove_var("BRZ_HOME");
-        }
-        let _ = std::env::set_current_dir(&self.old_cwd);
-        pyo3::Python::with_gil(|py| {
-            let os = py.import_bound("os").unwrap();
-            os.call_method1("chdir", (self.old_cwd.to_str().unwrap(),))
-                .unwrap();
-            if let Some(dir) = self.old_home.as_ref() {
-                os.call_method1("putenv", ("HOME", dir)).unwrap();
+        for (key, value) in self.old_env.iter() {
+            if let Some(value) = value {
+                std::env::set_var(key, value);
             } else {
-                os.call_method1("unsetenv", ("HOME",)).unwrap();
+                std::env::remove_var(key);
             }
-            if let Some(email) = self.old_email.as_ref() {
-                os.call_method1("putenv", ("BRZ_EMAIL", email)).unwrap();
-            } else {
-                os.call_method1("unsetenv", ("BRZ_EMAIL",)).unwrap();
-            }
-            if let Some(dir) = self.old_brz_home.as_ref() {
-                os.call_method1("putenv", ("BRZ_HOME", dir)).unwrap();
-            } else {
-                os.call_method1("unsetenv", ("BRZ_HOME",)).unwrap();
-            }
-        });
+            Python::with_gil(|py| {
+                let os = py.import_bound("os").unwrap();
+                if let Some(value) = value {
+                    os.call_method1("putenv", (key, value)).unwrap();
+                } else {
+                    os.call_method1("unsetenv", (key,)).unwrap();
+                }
+            });
+        }
+        let _ = std::env::set_current_dir(&self.old_cwd);
     }
 }