tests: improve robustness of color_format tests (#329)

Previously, these tests would have spurious failures when NO_COLOR or
FORCE_COLOR was set in the user's environment, since we weren't clearing
one variable before testing a value for the other one. The previous
version of the code also did not restore environment variable values on
panic, which could cause spurious failures in other tests after one test
fails.

tests/color_format.rs

@@ -3,6 +3,7 @@
 use lazy_static::lazy_static;
 use miette::{Diagnostic, MietteHandler, MietteHandlerOpts, ReportHandler, RgbColors};
 use regex::Regex;
+use std::ffi::OsString;
 use std::fmt::{self, Debug};
 use std::sync::Mutex;
 use thiserror::Error;
@@ -42,16 +43,29 @@ fn color_format(handler: MietteHandler) -> ColorFormat {
     }
 }
 
-/// Runs a function with an environment variable set to a specific value, then
-/// sets it back to it's original value once completed.
-fn with_env_var<F: FnOnce()>(var: &str, value: &str, body: F) {
-    let old_value = std::env::var_os(var);
-    std::env::set_var(var, value);
-    body();
-    if let Some(old_value) = old_value {
-        std::env::set_var(var, old_value);
-    } else {
-        std::env::remove_var(var);
+/// Store the current value of an environment variable on construction, and then
+/// restore that value when the guard is dropped.
+struct EnvVarGuard<'a> {
+    var: &'a str,
+    old_value: Option<OsString>,
+}
+
+impl EnvVarGuard<'_> {
+    fn new(var: &str) -> EnvVarGuard<'_> {
+        EnvVarGuard {
+            var,
+            old_value: std::env::var_os(var),
+        }
+    }
+}
+
+impl Drop for EnvVarGuard<'_> {
+    fn drop(&mut self) {
+        if let Some(old_value) = &self.old_value {
+            std::env::set_var(self.var, old_value);
+        } else {
+            std::env::remove_var(self.var);
+        }
     }
 }
 
@@ -72,22 +86,33 @@ fn check_colors<F: Fn(MietteHandlerOpts) -> MietteHandlerOpts>(
     //
     // Since environment variables are shared for the entire process, we need
     // to ensure that only one test that modifies these env vars runs at a time.
-    let guard = COLOR_ENV_VARS.lock().unwrap();
-
-    with_env_var("NO_COLOR", "1", || {
-        let handler = make_handler(MietteHandlerOpts::new()).build();
-        assert_eq!(color_format(handler), no_support);
-    });
-    with_env_var("FORCE_COLOR", "1", || {
-        let handler = make_handler(MietteHandlerOpts::new()).build();
-        assert_eq!(color_format(handler), ansi_support);
-    });
-    with_env_var("FORCE_COLOR", "3", || {
-        let handler = make_handler(MietteHandlerOpts::new()).build();
-        assert_eq!(color_format(handler), rgb_support);
-    });
-
-    drop(guard);
+    let lock = COLOR_ENV_VARS.lock().unwrap();
+
+    let guards = (
+        EnvVarGuard::new("NO_COLOR"),
+        EnvVarGuard::new("FORCE_COLOR"),
+    );
+    // Clear color environment variables that may be set outside of 'cargo test'
+    std::env::remove_var("NO_COLOR");
+    std::env::remove_var("FORCE_COLOR");
+
+    std::env::set_var("NO_COLOR", "1");
+    let handler = make_handler(MietteHandlerOpts::new()).build();
+    assert_eq!(color_format(handler), no_support);
+    std::env::remove_var("NO_COLOR");
+
+    std::env::set_var("FORCE_COLOR", "1");
+    let handler = make_handler(MietteHandlerOpts::new()).build();
+    assert_eq!(color_format(handler), ansi_support);
+    std::env::remove_var("FORCE_COLOR");
+
+    std::env::set_var("FORCE_COLOR", "3");
+    let handler = make_handler(MietteHandlerOpts::new()).build();
+    assert_eq!(color_format(handler), rgb_support);
+    std::env::remove_var("FORCE_COLOR");
+
+    drop(guards);
+    drop(lock);
 }
 
 #[test]