temp_env/
lib.rs

1#![deny(missing_docs)]
2//! This crate is for setting environment variables temporarily.
3//!
4//! It is useful for testing with different environment variables that should not interfere.
5//!
6//! # Examples
7//!
8//! ```rust
9//! temp_env::with_var("MY_ENV_VAR", Some("production"), || {
10//!     // Run some code where `MY_ENV_VAR` set to `"production"`.
11//! });
12//!
13//! temp_env::with_vars(
14//!     [
15//!         ("FIRST_VAR", Some("Hello")),
16//!         ("SECOND_VAR", Some("World!")),
17//!     ],
18//!     || {
19//!         // Run some code where `FIRST_VAR` is set to `"Hello"` and `SECOND_VAR` is set to
20//!         // `"World!"`.
21//!     }
22//! );
23//!
24//! temp_env::with_vars(
25//!     [
26//!         ("FIRST_VAR", Some("Hello")),
27//!         ("SECOND_VAR", None),
28//!     ],
29//!     || {
30//!         // Run some code where `FIRST_VAR` is set to `"Hello"` and `SECOND_VAR` is unset (even if
31//!         // it was set before)
32//!     }
33//! );
34//! ```
35//!
36//! It's possible the closure returns a value:
37//!
38//! ```rust
39//! let s = temp_env::with_var("INNER_ENV_VAR", Some("inner value"), || {
40//!      std::env::var("INNER_ENV_VAR").unwrap()
41//! });
42//! println!("{}", s);
43//! ```
44//!
45
46use std::collections::HashMap;
47use std::env;
48use std::ffi::{OsStr, OsString};
49use std::hash::Hash;
50
51use parking_lot::{ReentrantMutex, ReentrantMutexGuard};
52
53/// Make sure that the environment isn't modified concurrently.
54static SERIAL_TEST: ReentrantMutex<()> = ReentrantMutex::new(());
55
56/// Sets a single environment variable for the duration of the closure.
57///
58/// The previous value is restored when the closure completes or panics, before unwinding the
59/// panic.
60///
61/// If `value` is set to `None`, then the environment variable is unset.
62pub fn with_var<K, V, F, R>(key: K, value: Option<V>, closure: F) -> R
63where
64    K: AsRef<OsStr> + Clone + Eq + Hash,
65    V: AsRef<OsStr> + Clone,
66    F: FnOnce() -> R,
67{
68    with_vars([(key, value)], closure)
69}
70
71/// Unsets a single environment variable for the duration of the closure.
72///
73/// The previous value is restored when the closure completes or panics, before unwinding the
74/// panic.
75///
76/// This is a shorthand and identical to the following:
77/// ```rust
78/// temp_env::with_var("MY_ENV_VAR", None::<&str>, || {
79///     // Run some code where `MY_ENV_VAR` is unset.
80/// });
81/// ```
82pub fn with_var_unset<K, F, R>(key: K, closure: F) -> R
83where
84    K: AsRef<OsStr> + Clone + Eq + Hash,
85    F: FnOnce() -> R,
86{
87    with_var(key, None::<&str>, closure)
88}
89
90struct RestoreEnv<'a> {
91    env: HashMap<&'a OsStr, Option<OsString>>,
92    _guard: ReentrantMutexGuard<'a, ()>,
93}
94
95impl<'a> RestoreEnv<'a> {
96    /// Capture the given variables from the environment.
97    ///
98    /// `guard` holds a lock on the shared mutex for exclusive access to the environment, to make
99    /// sure that the environment gets restored while the lock is still held, i.e the current
100    /// thread still has exclusive access to the environment.
101    fn capture<I>(guard: ReentrantMutexGuard<'a, ()>, vars: I) -> Self
102    where
103        I: Iterator<Item = &'a OsStr> + 'a,
104    {
105        let env = vars.map(|v| (v, env::var_os(v))).collect();
106        Self { env, _guard: guard }
107    }
108}
109
110impl<'a> Drop for RestoreEnv<'a> {
111    fn drop(&mut self) {
112        for (var, value) in self.env.iter() {
113            update_env(var, value.as_ref().map(|v| v.as_os_str()));
114        }
115    }
116}
117
118/// Sets environment variables for the duration of the closure.
119///
120/// The previous values are restored when the closure completes or panics, before unwinding the
121/// panic.
122///
123/// If a `value` is set to `None`, then the environment variable is unset.
124///
125/// If the variable with the same name is set multiple times, the last one wins.
126pub fn with_vars<K, V, F, R>(kvs: impl AsRef<[(K, Option<V>)]>, closure: F) -> R
127where
128    K: AsRef<OsStr> + Clone + Eq + Hash,
129    V: AsRef<OsStr> + Clone,
130    F: FnOnce() -> R,
131{
132    let old_env = RestoreEnv::capture(
133        SERIAL_TEST.lock(),
134        kvs.as_ref().iter().map(|(k, _)| k.as_ref()),
135    );
136    for (key, value) in kvs.as_ref() {
137        update_env(key, value.as_ref());
138    }
139    let retval = closure();
140    drop(old_env);
141    retval
142}
143
144/// Unsets environment variables for the duration of the closure.
145///
146/// The previous values are restored when the closure completes or panics, before unwinding the
147/// panic.
148///
149/// This is a shorthand and identical to the following:
150/// ```rust
151/// temp_env::with_vars(
152///     [
153///         ("FIRST_VAR", None::<&str>),
154///         ("SECOND_VAR", None::<&str>),
155///     ],
156///     || {
157///         // Run some code where `FIRST_VAR` and `SECOND_VAR` are unset (even if
158///         // they were set before)
159///     }
160/// );
161/// ```
162pub fn with_vars_unset<K, F, R>(keys: impl AsRef<[K]>, closure: F) -> R
163where
164    K: AsRef<OsStr> + Clone + Eq + Hash,
165    F: FnOnce() -> R,
166{
167    let kvs = keys
168        .as_ref()
169        .iter()
170        .map(|key| (key, None::<&str>))
171        .collect::<Vec<_>>();
172    with_vars(kvs, closure)
173}
174
175fn update_env<K, V>(key: K, value: Option<V>)
176where
177    K: AsRef<OsStr>,
178    V: AsRef<OsStr>,
179{
180    match value {
181        Some(v) => env::set_var(key, v),
182        None => env::remove_var(key),
183    }
184}
185
186#[cfg(feature = "async_closure")]
187/// Does the same as [`with_vars`] but it allows to pass an async closures.
188///
189/// ```rust
190/// async fn check_var() {
191///     let v = std::env::var("MY_VAR").unwrap();
192///     assert_eq!(v, "ok".to_owned());
193/// }
194
195/// #[tokio::test]
196/// async fn test_async_closure() {
197///     crate::async_with_vars([("MY_VAR", Some("ok"))], check_var());
198/// }
199/// ```
200pub async fn async_with_vars<K, V, F, R>(kvs: impl AsRef<[(K, Option<V>)]>, closure: F) -> R
201where
202    K: AsRef<OsStr> + Clone + Eq + Hash,
203    V: AsRef<OsStr> + Clone,
204    F: std::future::Future<Output = R> + std::future::IntoFuture<Output = R>,
205{
206    let old_env = RestoreEnv::capture(
207        SERIAL_TEST.lock(),
208        kvs.as_ref().iter().map(|(k, _)| k.as_ref()),
209    );
210    for (key, value) in kvs.as_ref() {
211        update_env(key, value.as_ref());
212    }
213    let retval = closure.await;
214    drop(old_env);
215    retval
216}
217
218// Make sure that all tests use independent environment variables, so that they don't interfere if
219// run in parallel.
220#[cfg(test)]
221mod tests {
222    use std::env::VarError;
223    use std::{env, panic};
224
225    /// Test whether setting a variable is correctly undone.
226    #[test]
227    fn test_with_var_set() {
228        let hello_not_set = env::var("HELLO");
229        assert!(hello_not_set.is_err(), "`HELLO` must not be set.");
230
231        crate::with_var("HELLO", Some("world!"), || {
232            let hello_is_set = env::var("HELLO").unwrap();
233            assert_eq!(hello_is_set, "world!", "`HELLO` must be set to \"world!\".");
234        });
235
236        let hello_not_set_after = env::var("HELLO");
237        assert!(hello_not_set_after.is_err(), "`HELLO` must not be set.");
238    }
239
240    /// Test whether unsetting a variable is correctly undone.
241    #[test]
242    fn test_with_var_set_to_none() {
243        env::set_var("FOO", "bar");
244        let foo_is_set = env::var("FOO").unwrap();
245        assert_eq!(foo_is_set, "bar", "`FOO` must be set to \"bar\".");
246
247        crate::with_var("FOO", None::<&str>, || {
248            let foo_not_set = env::var("FOO");
249            assert!(foo_not_set.is_err(), "`FOO` must not be set.");
250        });
251
252        let foo_is_set_after = env::var("FOO").unwrap();
253        assert_eq!(foo_is_set_after, "bar", "`FOO` must be set to \"bar\".");
254    }
255
256    /// Test whether unsetting a variable through the shorthand is correctly undone.
257    #[test]
258    fn test_with_var_unset() {
259        env::set_var("BAR", "foo");
260        let foo_is_set = env::var("BAR").unwrap();
261        assert_eq!(foo_is_set, "foo", "`BAR` must be set to \"foo\".");
262
263        crate::with_var_unset("BAR", || {
264            let foo_not_set = env::var("BAR");
265            assert!(foo_not_set.is_err(), "`BAR` must not be set.");
266        });
267
268        let foo_is_set_after = env::var("BAR").unwrap();
269        assert_eq!(foo_is_set_after, "foo", "`BAR` must be set to \"foo\".");
270    }
271
272    /// Test whether overriding an existing variable is correctly undone.
273    #[test]
274    fn test_with_var_override() {
275        env::set_var("BLAH", "blub");
276        let blah_is_set = env::var("BLAH").unwrap();
277        assert_eq!(blah_is_set, "blub", "`BLAH` must be set to \"blah\".");
278
279        crate::with_var("BLAH", Some("new"), || {
280            let blah_is_set_new = env::var("BLAH").unwrap();
281            assert_eq!(blah_is_set_new, "new", "`BLAH` must be set to \"newb\".");
282        });
283
284        let blah_is_set_after = env::var("BLAH").unwrap();
285        assert_eq!(
286            blah_is_set_after, "blub",
287            "`BLAH` must be set to \"blubr\"."
288        );
289    }
290
291    /// Test whether overriding a variable is correctly undone in case of a panic.
292    #[test]
293    fn test_with_var_panic() {
294        env::set_var("PANIC", "panic");
295        let panic_is_set = env::var("PANIC").unwrap();
296        assert_eq!(panic_is_set, "panic", "`PANIC` must be set to \"panic\".");
297
298        let did_panic = panic::catch_unwind(|| {
299            crate::with_var("PANIC", Some("don't panic"), || {
300                let panic_is_set_new = env::var("PANIC").unwrap();
301                assert_eq!(
302                    panic_is_set_new, "don't panic",
303                    "`PANIC` must be set to \"don't panic\"."
304                );
305                panic!("abort this closure with a panic.");
306            });
307        });
308
309        assert!(did_panic.is_err(), "The closure must panic.");
310        let panic_is_set_after = env::var("PANIC").unwrap();
311        assert_eq!(
312            panic_is_set_after, "panic",
313            "`PANIC` must be set to \"panic\"."
314        );
315    }
316
317    /// Test whether setting multiple variable is correctly undone.
318    #[test]
319    fn test_with_vars_set() {
320        let one_not_set = env::var("ONE");
321        assert!(one_not_set.is_err(), "`ONE` must not be set.");
322        let two_not_set = env::var("TWO");
323        assert!(two_not_set.is_err(), "`TWO` must not be set.");
324
325        crate::with_vars([("ONE", Some("1")), ("TWO", Some("2"))], || {
326            let one_is_set = env::var("ONE").unwrap();
327            assert_eq!(one_is_set, "1", "`ONE` must be set to \"1\".");
328            let two_is_set = env::var("TWO").unwrap();
329            assert_eq!(two_is_set, "2", "`TWO` must be set to \"2\".");
330        });
331
332        let one_not_set_after = env::var("ONE");
333        assert!(one_not_set_after.is_err(), "`ONE` must not be set.");
334        let two_not_set_after = env::var("TWO");
335        assert!(two_not_set_after.is_err(), "`TWO` must not be set.");
336    }
337
338    /// Test whether setting multiple variable is returns result.
339    #[test]
340    fn test_with_vars_set_returning() {
341        let one_not_set = env::var("ONE");
342        assert!(one_not_set.is_err(), "`ONE` must not be set.");
343        let two_not_set = env::var("TWO");
344        assert!(two_not_set.is_err(), "`TWO` must not be set.");
345
346        let r = crate::with_vars([("ONE", Some("1")), ("TWO", Some("2"))], || {
347            let one_is_set = env::var("ONE").unwrap();
348            let two_is_set = env::var("TWO").unwrap();
349            (one_is_set, two_is_set)
350        });
351
352        let (one_from_closure, two_from_closure) = r;
353
354        assert_eq!(one_from_closure, "1", "`ONE` had to be set to \"1\".");
355        assert_eq!(two_from_closure, "2", "`TWO` had to be set to \"2\".");
356
357        let one_not_set_after = env::var("ONE");
358        assert!(one_not_set_after.is_err(), "`ONE` must not be set.");
359        let two_not_set_after = env::var("TWO");
360        assert!(two_not_set_after.is_err(), "`TWO` must not be set.");
361    }
362
363    /// Test whether unsetting multiple variables is correctly undone.
364    #[test]
365    fn test_with_vars_unset() {
366        env::set_var("SET_TO_BE_UNSET", "val");
367        env::remove_var("UNSET_TO_BE_UNSET");
368        // Check test preconditions
369        assert_eq!(env::var("SET_TO_BE_UNSET"), Ok("val".to_string()));
370        assert_eq!(env::var("UNSET_TO_BE_UNSET"), Err(VarError::NotPresent));
371
372        crate::with_vars_unset(["SET_TO_BE_UNSET", "UNSET_TO_BE_UNSET"], || {
373            assert_eq!(env::var("SET_TO_BE_UNSET"), Err(VarError::NotPresent));
374            assert_eq!(env::var("UNSET_TO_BE_UNSET"), Err(VarError::NotPresent));
375        });
376        assert_eq!(env::var("SET_TO_BE_UNSET"), Ok("val".to_string()));
377        assert_eq!(env::var("UNSET_TO_BE_UNSET"), Err(VarError::NotPresent));
378    }
379
380    /// Test whether unsetting one of the variable is correctly undone.
381    #[test]
382    fn test_with_vars_partially_unset() {
383        let to_be_set_not_set = env::var("TO_BE_SET");
384        assert!(to_be_set_not_set.is_err(), "`TO_BE_SET` must not be set.");
385        env::set_var("TO_BE_UNSET", "unset");
386        let to_be_unset_is_set = env::var("TO_BE_UNSET").unwrap();
387        assert_eq!(
388            to_be_unset_is_set, "unset",
389            "`TO_BE_UNSET` must be set to \"unset\"."
390        );
391
392        crate::with_vars(
393            [("TO_BE_SET", Some("set")), ("TO_BE_UNSET", None::<&str>)],
394            || {
395                let to_be_set_is_set = env::var("TO_BE_SET").unwrap();
396                assert_eq!(
397                    to_be_set_is_set, "set",
398                    "`TO_BE_SET` must be set to \"set\"."
399                );
400                let to_be_unset_not_set = env::var("TO_BE_UNSET");
401                assert!(
402                    to_be_unset_not_set.is_err(),
403                    "`TO_BE_UNSET` must not be set."
404                );
405            },
406        );
407
408        let to_be_set_not_set_after = env::var("TO_BE_SET");
409        assert!(
410            to_be_set_not_set_after.is_err(),
411            "`TO_BE_SET` must not be set."
412        );
413        let to_be_unset_is_set_after = env::var("TO_BE_UNSET").unwrap();
414        assert_eq!(
415            to_be_unset_is_set_after, "unset",
416            "`TO_BE_UNSET` must be set to \"unset\"."
417        );
418    }
419
420    /// Test whether overriding existing variables is correctly undone.
421    #[test]
422    fn test_with_vars_override() {
423        env::set_var("DOIT", "doit");
424        let doit_is_set = env::var("DOIT").unwrap();
425        assert_eq!(doit_is_set, "doit", "`DOIT` must be set to \"doit\".");
426        env::set_var("NOW", "now");
427        let now_is_set = env::var("NOW").unwrap();
428        assert_eq!(now_is_set, "now", "`NOW` must be set to \"now\".");
429
430        crate::with_vars([("DOIT", Some("other")), ("NOW", Some("value"))], || {
431            let doit_is_set_new = env::var("DOIT").unwrap();
432            assert_eq!(doit_is_set_new, "other", "`DOIT` must be set to \"other\".");
433            let now_is_set_new = env::var("NOW").unwrap();
434            assert_eq!(now_is_set_new, "value", "`NOW` must be set to \"value\".");
435        });
436
437        let doit_is_set_after = env::var("DOIT").unwrap();
438        assert_eq!(doit_is_set_after, "doit", "`DOIT` must be set to \"doit\".");
439        let now_is_set_after = env::var("NOW").unwrap();
440        assert_eq!(now_is_set_after, "now", "`NOW` must be set to \"now\".");
441    }
442
443    /// Test that setting the same variables twice, the latter one is used.
444    #[test]
445    fn test_with_vars_same_vars() {
446        let override_not_set = env::var("OVERRIDE");
447        assert!(override_not_set.is_err(), "`OVERRIDE` must not be set.");
448
449        crate::with_vars(
450            [
451                ("OVERRIDE", Some("initial")),
452                ("OVERRIDE", Some("override")),
453            ],
454            || {
455                let override_is_set = env::var("OVERRIDE").unwrap();
456                assert_eq!(
457                    override_is_set, "override",
458                    "`OVERRIDE` must be set to \"override\"."
459                );
460            },
461        );
462
463        let override_not_set_after = env::var("OVERRIDE");
464        assert!(
465            override_not_set_after.is_err(),
466            "`OVERRIDE` must not be set."
467        );
468    }
469
470    /// Test that unsetting and setting the same variable leads to the variable being set.
471    #[test]
472    fn test_with_vars_unset_set() {
473        env::set_var("MY_VAR", "my_var");
474        let my_var_is_set = env::var("MY_VAR").unwrap();
475        assert_eq!(
476            my_var_is_set, "my_var",
477            "`MY_VAR` must be set to \"my_var`\"."
478        );
479
480        crate::with_vars(
481            [("MY_VAR", None::<&str>), ("MY_VAR", Some("new value"))],
482            || {
483                let my_var_is_set_new = env::var("MY_VAR").unwrap();
484                assert_eq!(
485                    my_var_is_set_new, "new value",
486                    "`MY_VAR` must be set to \"new value\"."
487                );
488            },
489        );
490
491        let my_var_is_set_after = env::var("MY_VAR").unwrap();
492        assert_eq!(
493            my_var_is_set_after, "my_var",
494            "`MY_VAR` must be set to \"my_var\"."
495        );
496    }
497
498    /// Test that setting and unsetting the same variable leads to the variable being unset.
499    #[test]
500    fn test_with_vars_set_unset() {
501        let not_my_var_not_set = env::var("NOT_MY_VAR");
502        assert!(not_my_var_not_set.is_err(), "`NOT_MY_VAR` must not be set.");
503
504        crate::with_vars(
505            [
506                ("NOT_MY_VAR", Some("it is set")),
507                ("NOT_MY_VAR", None::<&str>),
508            ],
509            || {
510                let not_my_var_not_set_new = env::var("NOT_MY_VAR");
511                assert!(
512                    not_my_var_not_set_new.is_err(),
513                    "`NOT_MY_VAR` must not be set."
514                );
515            },
516        );
517
518        let not_my_var_not_set_after = env::var("NOT_MY_VAR");
519        assert!(
520            not_my_var_not_set_after.is_err(),
521            "`NOT_MY_VAR` must not be set."
522        );
523    }
524
525    #[test]
526    fn test_with_nested_set() {
527        crate::with_var("MY_VAR_1", Some("1"), || {
528            crate::with_var("MY_VAR_2", Some("2"), || {
529                assert_eq!(env::var("MY_VAR_1").unwrap(), "1");
530                assert_eq!(env::var("MY_VAR_2").unwrap(), "2");
531            })
532        });
533
534        assert!(env::var("MY_VAR_1").is_err());
535        assert!(env::var("MY_VAR_2").is_err());
536    }
537
538    #[test]
539    fn test_fn_once() {
540        let value = String::from("Hello, ");
541        let value = crate::with_var("WORLD", Some("world!"), || {
542            value + &env::var("WORLD").unwrap()
543        });
544        assert_eq!(value, "Hello, world!");
545    }
546
547    #[cfg(feature = "async_closure")]
548    async fn check_var() {
549        let v = std::env::var("MY_VAR").unwrap();
550        assert_eq!(v, "ok".to_owned());
551    }
552
553    #[cfg(feature = "async_closure")]
554    #[tokio::test]
555    async fn test_async_closure() {
556        crate::async_with_vars([("MY_VAR", Some("ok"))], check_var()).await;
557        let f = async {
558            let v = std::env::var("MY_VAR").unwrap();
559            assert_eq!(v, "ok".to_owned());
560        };
561        crate::async_with_vars([("MY_VAR", Some("ok"))], f).await;
562    }
563
564    #[cfg(feature = "async_closure")]
565    #[tokio::test(flavor = "multi_thread")]
566    async fn test_async_closure_calls_closure() {
567        let (tx, rx) = tokio::sync::oneshot::channel();
568        let f = async {
569            tx.send(std::env::var("MY_VAR")).unwrap();
570        };
571        crate::async_with_vars([("MY_VAR", Some("ok"))], f).await;
572        let value = rx.await.unwrap().unwrap();
573        assert_eq!(value, "ok".to_owned());
574    }
575
576    #[cfg(feature = "async_closure")]
577    #[tokio::test(flavor = "multi_thread")]
578    async fn test_async_with_vars_set_returning() {
579        let one_not_set = env::var("ONE");
580        assert!(one_not_set.is_err(), "`ONE` must not be set.");
581        let two_not_set = env::var("TWO");
582        assert!(two_not_set.is_err(), "`TWO` must not be set.");
583
584        let r = crate::async_with_vars([("ONE", Some("1")), ("TWO", Some("2"))], async {
585            let one_is_set = env::var("ONE").unwrap();
586            let two_is_set = env::var("TWO").unwrap();
587            (one_is_set, two_is_set)
588        })
589        .await;
590
591        let (one_from_closure, two_from_closure) = r;
592
593        assert_eq!(one_from_closure, "1", "`ONE` had to be set to \"1\".");
594        assert_eq!(two_from_closure, "2", "`TWO` had to be set to \"2\".");
595
596        let one_not_set_after = env::var("ONE");
597        assert!(one_not_set_after.is_err(), "`ONE` must not be set.");
598        let two_not_set_after = env::var("TWO");
599        assert!(two_not_set_after.is_err(), "`TWO` must not be set.");
600    }
601}
temp_env/lib.rs

temp_env/
lib.rs
