use once_cell::sync::OnceCell; use std::collections::HashMap; use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex, PoisonError}; use std::{env, fs, io}; use tempfile::{tempdir, TempDir}; use super::*; /// Env var convenience type. type EnvMap = HashMap; /// Initialized in [`get_env_locker`] static ENV_LOCKER: OnceCell>> = OnceCell::new(); /// A test environment. /// /// Will create a new temporary directory. Use its builder methods to configure /// the directory structure, preset variables, envfile name and contents, and /// the working directory to run the test from. /// /// Creation methods: /// - [`TestEnv::init`]: blank environment (no envfile) /// - [`TestEnv::init_with_envfile`]: blank environment with an envfile /// - [`TestEnv::default`]: default testing environment (1 existing var and 2 /// set in a `.env` file) #[derive(Debug)] pub struct TestEnv { temp_dir: TempDir, work_dir: PathBuf, env_vars: Vec, envfile_contents: Option, envfile_path: PathBuf, } /// Simple key value struct for representing environment variables #[derive(Debug, Clone)] pub struct KeyVal { key: String, value: String, } /// Run a test closure within a test environment. /// /// Resets the environment variables, loads the [`TestEnv`], then runs the test /// closure. Ensures only one thread has access to the process environment. pub fn test_in_env(test_env: TestEnv, test: F) where F: FnOnce(), { let locker = get_env_locker(); // ignore a poisoned mutex // we expect some tests may panic to indicate a failure let original_env = locker.lock().unwrap_or_else(PoisonError::into_inner); // we reset the environment anyway upon acquiring the lock reset_env(&original_env); create_env(&test_env); test(); // drop the lock and the `TestEnv` - should delete the tempdir } /// Run a test closure within the default test environment. /// /// Resets the environment variables, creates the default [`TestEnv`], then runs /// the test closure. Ensures only one thread has access to the process /// environment. /// /// The default testing environment sets an existing environment variable /// `TEST_EXISTING_KEY`, which is set to `from_env`. It also creates a `.env` /// file with the two lines: /// /// ```ini /// TESTKEY=test_val /// TEST_EXISTING_KEY=from_file /// ``` /// /// Notice that file has the potential to override `TEST_EXISTING_KEY` depending /// on the what's being tested. pub fn test_in_default_env(test: F) where F: FnOnce(), { let test_env = TestEnv::default(); test_in_env(test_env, test); } impl TestEnv { /// Blank testing environment in a new temporary directory. /// /// No envfile_contents or pre-existing variables to set. The envfile_name /// is set to `.env` but wont be written until its contents is set. The /// working directory is the created temporary directory. pub fn init() -> Self { let tempdir = tempdir().expect("create tempdir"); let work_dir = tempdir.path().to_owned(); let envfile_path = work_dir.join(".env"); Self { temp_dir: tempdir, work_dir, env_vars: Default::default(), envfile_contents: None, envfile_path, } } /// Testing environment with custom envfile_contents. /// /// No pre-existing env_vars set. The envfile_name is set to `.env`. The /// working directory is the created temporary directory. pub fn init_with_envfile(contents: impl ToString) -> Self { let mut test_env = Self::init(); test_env.set_envfile_contents(contents); test_env } /// Change the name of the default `.env` file. /// /// It will still be placed in the root temporary directory. If you need to /// put the envfile in a different directory, use /// [`set_envfile_path`](TestEnv::set_envfile_path) instead. pub fn set_envfile_name(&mut self, name: impl AsRef) -> &mut Self { self.envfile_path = self.temp_path().join(name); self } /// Change the absolute path to the envfile. pub fn set_envfile_path(&mut self, path: PathBuf) -> &mut Self { self.envfile_path = path; self } /// Specify the contents of the envfile. /// /// If this is the only change to the [`TestEnv`] being made, use /// [`new_with_envfile`](TestEnv::new_with_envfile). /// /// Setting it to an empty string will cause an empty envfile to be created pub fn set_envfile_contents(&mut self, contents: impl ToString) -> &mut Self { self.envfile_contents = Some(contents.to_string()); self } /// Set the working directory the test will run from. /// /// The default is the created temporary directory. This method is useful if /// you wish to run a test from a subdirectory or somewhere else. pub fn set_work_dir(&mut self, path: PathBuf) -> &mut Self { self.work_dir = path; self } /// Add an individual environment variable. /// /// This adds more pre-existing environment variables to the process before /// any tests are run. pub fn add_env_var(&mut self, key: impl ToString, value: impl ToString) -> &mut Self { self.env_vars.push(KeyVal { key: key.to_string(), value: value.to_string(), }); self } /// Set the pre-existing environment variables. /// /// These variables will get added to the process' environment before the /// test is run. This overrides any previous env vars added to the /// [`TestEnv`]. /// /// If you wish to just use a slice of tuples, use /// [`set_env_vars_tuple`](TestEnv::set_env_vars_tuple) instead. pub fn set_env_vars(&mut self, env_vars: Vec) -> &mut Self { self.env_vars = env_vars; self } /// Set the pre-existing environment variables using [`str`] tuples. /// /// These variables will get added to the process' environment before the /// test is run. This overrides any previous env vars added to the /// [`TestEnv`]. /// /// If you wish to add an owned `Vec` instead of `str` tuples, use /// [`set_env_vars`](TestEnv::set_env_vars) instead. pub fn set_env_vars_tuple(&mut self, env_vars: &[(&str, &str)]) -> &mut Self { self.env_vars = env_vars .iter() .map(|(key, value)| KeyVal { key: key.to_string(), value: value.to_string(), }) .collect(); self } /// Create a child folder within the temporary directory. /// /// This will not change the working directory the test is run in, or where /// the envfile is created. /// /// Will create parent directories if they are missing. pub fn add_child_dir_all(&self, rel_path: impl AsRef) -> PathBuf { let rel_path = rel_path.as_ref(); let child_dir = self.temp_path().join(rel_path); if let Err(err) = fs::create_dir_all(&child_dir) { panic!( "unable to create child directory: `{}` in `{}`: {}", self.temp_path().display(), rel_path.display(), err ); } child_dir } /// Get a reference to the path of the temporary directory. pub fn temp_path(&self) -> &Path { self.temp_dir.path() } /// Get a reference to the working directory the test will be run from. pub fn work_dir(&self) -> &Path { &self.work_dir } /// Get a reference to environnement variables that will be set **before** /// the test. pub fn env_vars(&self) -> &[KeyVal] { &self.env_vars } /// Get a reference to the string that will be placed in the envfile. /// /// If `None` is returned, an envfile will not be created pub fn envfile_contents(&self) -> Option<&str> { self.envfile_contents.as_deref() } /// Get a reference to the path of the envfile. pub fn envfile_path(&self) -> &Path { &self.envfile_path } } impl Default for TestEnv { fn default() -> Self { let temp_dir = tempdir().expect("create tempdir"); let work_dir = temp_dir.path().to_owned(); let env_vars = vec![KeyVal { key: TEST_EXISTING_KEY.into(), value: TEST_EXISTING_VALUE.into(), }]; let envfile_contents = Some(create_default_envfile()); let envfile_path = work_dir.join(".env"); Self { temp_dir, work_dir, env_vars, envfile_contents, envfile_path, } } } impl From<(&str, &str)> for KeyVal { fn from(kv: (&str, &str)) -> Self { let (key, value) = kv; Self { key: key.to_string(), value: value.to_string(), } } } impl From<(String, String)> for KeyVal { fn from(kv: (String, String)) -> Self { let (key, value) = kv; Self { key, value } } } /// Get a guarded copy of the original process' env vars. fn get_env_locker() -> Arc> { Arc::clone(ENV_LOCKER.get_or_init(|| { let map: EnvMap = env::vars().collect(); Arc::new(Mutex::new(map)) })) } /// Reset the process' env vars back to what was in `original_env`. fn reset_env(original_env: &EnvMap) { // remove keys if they weren't in the original environment env::vars() .filter(|(key, _)| !original_env.contains_key(key)) .for_each(|(key, _)| env::remove_var(key)); // ensure original keys have their original values original_env .iter() .for_each(|(key, value)| env::set_var(key, value)); } /// Create an environment to run tests in. /// /// Writes the envfile, sets the working directory, and sets environment vars. fn create_env(test_env: &TestEnv) { // only create the envfile if its contents has been set if let Some(contents) = test_env.envfile_contents() { create_envfile(&test_env.envfile_path, contents); } env::set_current_dir(&test_env.work_dir).expect("setting working directory"); for KeyVal { key, value } in &test_env.env_vars { env::set_var(key, value) } } /// Create an envfile for use in tests. fn create_envfile(path: &Path, contents: &str) { if path.exists() { panic!("envfile `{}` already exists", path.display()) } // inner function to group together io::Results fn create_env_file_inner(path: &Path, contents: &str) -> io::Result<()> { let mut file = fs::File::create(path)?; file.write_all(contents.as_bytes())?; file.sync_all() } // call inner function if let Err(err) = create_env_file_inner(path, contents) { // handle any io::Result::Err panic!("error creating envfile `{}`: {}", path.display(), err); } }