tests: allow setting env var for all commands in TestEnvironment

Sometimes it's useful to have an environment variable set for all
commands in a test. This patch lets you do that by adding environment
variables to the `TestEnvironment` itself. These will then be set on
all subsequent commands.

tests/common/mod.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::cell::RefCell;
+use std::collections::HashMap;
 use std::io::Write;
 use std::path::{Path, PathBuf};
 
@@ -23,6 +24,7 @@ pub struct TestEnvironment {
     env_root: PathBuf,
     home_dir: PathBuf,
     config_path: PathBuf,
+    env_vars: HashMap<String, String>,
     command_number: RefCell<i64>,
 }
 
@@ -34,11 +36,13 @@ impl Default for TestEnvironment {
         std::fs::create_dir(&home_dir).unwrap();
         let config_path = env_root.join("config.toml");
         std::fs::write(&config_path, b"").unwrap();
+        let env_vars = HashMap::new();
         Self {
             _temp_dir: tmp_dir,
             env_root,
             home_dir,
             config_path,
+            env_vars,
             command_number: RefCell::new(0),
         }
     }
@@ -50,6 +54,9 @@ impl TestEnvironment {
         cmd.current_dir(current_dir);
         cmd.args(args);
         cmd.env_clear();
+        for (key, value) in &self.env_vars {
+            cmd.env(key, value);
+        }
         cmd.env("RUST_BACKTRACE", "1");
         cmd.env("HOME", self.home_dir.to_str().unwrap());
         let timestamp = chrono::DateTime::parse_from_rfc3339("2001-02-03T04:05:06+07:00").unwrap();
@@ -89,6 +96,10 @@ impl TestEnvironment {
         config_file.write_all(content).unwrap();
         config_file.flush().unwrap();
     }
+
+    pub fn add_env_var(&mut self, key: &str, val: &str) {
+        self.env_vars.insert(key.to_string(), val.to_string());
+    }
 }
 
 pub fn get_stdout_string(assert: &assert_cmd::assert::Assert) -> String {

tests/test_global_opts.rs

@@ -71,7 +71,7 @@ fn test_repo_arg_with_git_clone() {
 
 #[test]
 fn test_color_config() {
-    let test_env = TestEnvironment::default();
+    let mut test_env = TestEnvironment::default();
 
     // Test that color is used if it's requested in the config file
     test_env.write_config(
@@ -88,12 +88,9 @@ color="always""#,
     "###);
 
     // Test that NO_COLOR overrides the request for color in the config file
-    let assert = test_env
-        .jj_cmd(&repo_path, &["log", "-T", "commit_id"])
-        .env("NO_COLOR", "")
-        .assert()
-        .success();
-    insta::assert_snapshot!(get_stdout_string(&assert), @r###"
+    test_env.add_env_var("NO_COLOR", "");
+    let stdout = test_env.jj_cmd_success(&repo_path, &["log", "-T", "commit_id"]);
+    insta::assert_snapshot!(stdout, @r###"
     @ 230dd059e1b059aefc0da06a2e5a7dbf22362f22
     o 0000000000000000000000000000000000000000
     "###);