Refactor test registry state to be more type safe

Author: djc

src/cli/self_update/test.rs

@@ -6,12 +6,13 @@ use std::{io, sync::Mutex};
 #[cfg(windows)]
 use winreg::{
     enums::{HKEY_CURRENT_USER, KEY_READ, KEY_WRITE},
+    types::{FromRegValue, ToRegValue},
     RegKey, RegValue,
 };
 
 #[cfg(windows)]
 pub fn with_saved_path(f: &mut dyn FnMut()) {
-    with_saved_reg_value(&RegKey::predef(HKEY_CURRENT_USER), "Environment", "PATH", f)
+    with_saved_reg_value(&USER_PATH, f)
 }
 
 #[cfg(unix)]
@@ -21,41 +22,62 @@ pub fn with_saved_path(f: &mut dyn FnMut()) {
 
 #[cfg(windows)]
 pub fn get_path() -> io::Result<Option<RegValue>> {
-    get_reg_value(&RegKey::predef(HKEY_CURRENT_USER), "Environment", "PATH")
+    USER_PATH.get()
 }
 
 #[cfg(windows)]
-pub fn with_saved_reg_value(root: &RegKey, subkey: &str, name: &str, f: &mut dyn FnMut()) {
+pub fn with_saved_reg_value(id: &RegistryValueId, f: &mut dyn FnMut()) {
     // Lock protects concurrent mutation of registry
     static LOCK: Mutex<()> = Mutex::new(());
     let _g = LOCK.lock();
 
     // Save and restore the global state here to keep from trashing things.
-    let saved_state = get_reg_value(root, subkey, name)
+    let saved_state = id
+        .get()
         .expect("Error getting global state: Better abort to avoid trashing it");
-    let _g = scopeguard::guard(saved_state, |p| restore_reg_value(root, subkey, name, p));
+    let _g = scopeguard::guard(saved_state, |p| id.set(p.as_ref()).unwrap());
 
     f();
 }
 
 #[cfg(windows)]
-fn get_reg_value(root: &RegKey, subkey: &str, name: &str) -> io::Result<Option<RegValue>> {
-    let subkey = root.open_subkey_with_flags(subkey, KEY_READ | KEY_WRITE)?;
-    match subkey.get_raw_value(name) {
-        Ok(val) => Ok(Some(val)),
-        Err(ref e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
-        Err(e) => Err(e),
-    }
+const USER_PATH: RegistryValueId = RegistryValueId {
+    sub_key: "Environment",
+    value_name: "PATH",
+};
+
+#[cfg(windows)]
+pub struct RegistryValueId {
+    pub sub_key: &'static str,
+    pub value_name: &'static str,
 }
 
 #[cfg(windows)]
-fn restore_reg_value(root: &RegKey, subkey: &str, name: &str, p: Option<RegValue>) {
-    let subkey = root
-        .open_subkey_with_flags(subkey, KEY_READ | KEY_WRITE)
-        .unwrap();
-    if let Some(p) = p.as_ref() {
-        subkey.set_raw_value(name, p).unwrap();
-    } else {
-        let _ = subkey.delete_value(name);
+impl RegistryValueId {
+    pub fn get_value<T: FromRegValue>(&self) -> io::Result<Option<T>> {
+        self.get()?.map(|v| T::from_reg_value(&v)).transpose()
+    }
+
+    pub fn get(&self) -> io::Result<Option<RegValue>> {
+        let sub_key = RegKey::predef(HKEY_CURRENT_USER)
+            .open_subkey_with_flags(self.sub_key, KEY_READ | KEY_WRITE)?;
+        match sub_key.get_raw_value(self.value_name) {
+            Ok(val) => Ok(Some(val)),
+            Err(ref e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
+            Err(e) => Err(e),
+        }
+    }
+
+    pub fn set_value(&self, new: Option<impl ToRegValue>) -> io::Result<()> {
+        self.set(new.map(|s| s.to_reg_value()).as_ref())
+    }
+
+    pub fn set(&self, new: Option<&RegValue>) -> io::Result<()> {
+        let sub_key = RegKey::predef(HKEY_CURRENT_USER)
+            .open_subkey_with_flags(self.sub_key, KEY_READ | KEY_WRITE)?;
+        match new {
+            Some(new) => sub_key.set_raw_value(self.value_name, new),
+            None => sub_key.delete_value(self.value_name),
+        }
     }
 }

src/test.rs

@@ -20,7 +20,7 @@ use crate::currentprocess::TestProcess;
 use crate::dist::TargetTriple;
 
 #[cfg(windows)]
-pub use crate::cli::self_update::test::{get_path, with_saved_reg_value};
+pub use crate::cli::self_update::test::{get_path, with_saved_reg_value, RegistryValueId};
 
 // Things that can have environment variables applied to them.
 pub trait Env {

tests/suite/cli_self_upd.rs

@@ -19,12 +19,11 @@ use rustup::test::{
     },
     this_host_triple, with_saved_path,
 };
+#[cfg(windows)]
+use rustup::test::{with_saved_reg_value, RegistryValueId};
 use rustup::utils::{raw, utils};
 use rustup::{for_host, DUP_TOOLS, TOOLS};
 
-#[cfg(windows)]
-use rustup::test::with_saved_reg_value;
-
 const TEST_VERSION: &str = "1.1.1";
 
 /// Empty dist server, rustup installed with no toolchain
@@ -315,36 +314,31 @@ info: downloading self-update
 #[test]
 #[cfg(windows)]
 fn update_overwrites_programs_display_version() {
-    let root = &winreg::RegKey::predef(winreg::enums::HKEY_CURRENT_USER);
-    let key = r"Software\Microsoft\Windows\CurrentVersion\Uninstall\Rustup";
-    let name = "DisplayVersion";
-
     const PLACEHOLDER_VERSION: &str = "9.999.99";
     let version = env!("CARGO_PKG_VERSION");
 
     let mut cx = SelfUpdateTestContext::new(TEST_VERSION);
-    with_saved_reg_value(root, key, name, &mut || {
+    with_saved_reg_value(&USER_RUSTUP_VERSION, &mut || {
         cx.config
             .expect_ok(&["rustup-init", "-y", "--no-modify-path"]);
 
-        root.create_subkey(key)
-            .unwrap()
-            .0
-            .set_value(name, &PLACEHOLDER_VERSION)
+        USER_RUSTUP_VERSION
+            .set_value(Some(PLACEHOLDER_VERSION))
             .unwrap();
-
         cx.config.expect_ok(&["rustup", "self", "update"]);
-
         assert_eq!(
-            root.open_subkey(key)
-                .unwrap()
-                .get_value::<String, _>(name)
-                .unwrap(),
+            USER_RUSTUP_VERSION.get_value::<String>().unwrap().unwrap(),
             version,
         );
     });
 }
 
+#[cfg(windows)]
+const USER_RUSTUP_VERSION: RegistryValueId = RegistryValueId {
+    sub_key: r"Software\Microsoft\Windows\CurrentVersion\Uninstall\Rustup",
+    value_name: "DisplayVersion",
+};
+
 #[test]
 fn update_but_not_installed() {
     let cx = SelfUpdateTestContext::new(TEST_VERSION);