Move with_update_server() into CliTestContext

djc · djc · commit 29287aa23beb · 2024-06-22T16:58:35.000Z
diff --git a/src/test/mock/clitools.rs b/src/test/mock/clitools.rs
@@ -366,6 +366,43 @@ impl CliTestContext {
         self.config.distdir = Some(CONST_TEST_STATE.dist_server_for(scenario).unwrap());
         DistDirGuard { inner: self }
     }
+
+    pub fn with_update_server(&mut self, version: &str) -> UpdateServerGuard {
+        let self_dist_tmp = tempfile::Builder::new()
+            .prefix("self_dist")
+            .tempdir()
+            .unwrap();
+        let self_dist = self_dist_tmp.path();
+
+        let root_url = create_local_update_server(self_dist, &self.config.exedir, version);
+        let trip = this_host_triple();
+        let dist_dir = self_dist.join(format!("archive/{version}/{trip}"));
+        let dist_exe = dist_dir.join(format!("rustup-init{EXE_SUFFIX}"));
+        let dist_tmp = dist_dir.join("rustup-init-tmp");
+
+        // Modify the exe so it hashes different
+        // 1) move out of the way the file
+        fs::rename(&dist_exe, &dist_tmp).unwrap();
+        // 2) copy it
+        fs::copy(dist_tmp, &dist_exe).unwrap();
+        // modify it
+        let mut dest_file = fs::OpenOptions::new()
+            .append(true)
+            .create(true)
+            .open(dist_exe)
+            .unwrap();
+        writeln!(dest_file).unwrap();
+
+        self.config.rustup_update_root = Some(root_url);
+        UpdateServerGuard {
+            _self_dist: self_dist_tmp,
+        }
+    }
+}
+
+#[must_use]
+pub struct UpdateServerGuard {
+    _self_dist: TempDir,
 }
 
 impl From<Scenario> for CliTestContext {
@@ -384,10 +421,7 @@ impl From<Scenario> for CliTestContext {
             config.distdir = Some(config.test_dist_dir.path().to_path_buf());
         }
 
-        Self {
-            config,
-            _test_dir,
-        }
+        Self { config, _test_dir }
     }
 }
 
@@ -432,37 +466,6 @@ fn create_local_update_server(self_dist: &Path, exedir: &Path, version: &str) ->
     root_url
 }
 
-pub fn with_update_server(config: &mut Config, version: &str, f: &dyn Fn(&mut Config)) {
-    let self_dist_tmp = tempfile::Builder::new()
-        .prefix("self_dist")
-        .tempdir()
-        .unwrap();
-    let self_dist = self_dist_tmp.path();
-
-    let root_url = create_local_update_server(self_dist, &config.exedir, version);
-    let trip = this_host_triple();
-    let dist_dir = self_dist.join(format!("archive/{version}/{trip}"));
-    let dist_exe = dist_dir.join(format!("rustup-init{EXE_SUFFIX}"));
-    let dist_tmp = dist_dir.join("rustup-init-tmp");
-
-    // Modify the exe so it hashes different
-    // 1) move out of the way the file
-    fs::rename(&dist_exe, &dist_tmp).unwrap();
-    // 2) copy it
-    fs::copy(dist_tmp, &dist_exe).unwrap();
-    // modify it
-    let mut dest_file = fs::OpenOptions::new()
-        .append(true)
-        .create(true)
-        .open(dist_exe)
-        .unwrap();
-    writeln!(dest_file).unwrap();
-
-    config.rustup_update_root = Some(root_url);
-    f(config);
-    config.rustup_update_root = None;
-}
-
 pub fn output_release_file(dist_dir: &Path, schema: &str, version: &str) {
     let contents = format!(
         r#"
diff --git a/tests/suite/cli_exact.rs b/tests/suite/cli_exact.rs
@@ -3,7 +3,7 @@
 
 use rustup::for_host;
 use rustup::test::{
-    mock::clitools::{self, set_current_dist_date, with_update_server, CliTestContext, Scenario},
+    mock::clitools::{self, set_current_dist_date, CliTestContext, Scenario},
     this_host_triple,
 };
 
@@ -39,25 +39,27 @@ info: default toolchain set to 'nightly-{0}'
 fn update_once_and_check_self_update() {
     let test_version = "2.0.0";
     let mut cx = CliTestContext::from(Scenario::SimpleV2);
-    with_update_server(&mut cx.config, test_version, &|config| {
-        config.expect_ok(&["rustup-init", "-y", "--no-modify-path"]);
-        config.expect_ok(&["rustup", "set", "auto-self-update", "check-only"]);
-        let current_version = env!("CARGO_PKG_VERSION");
+    let _dist_guard = cx.with_update_server(test_version);
+    cx.config
+        .expect_ok(&["rustup-init", "-y", "--no-modify-path"]);
+    cx.config
+        .expect_ok(&["rustup", "set", "auto-self-update", "check-only"]);
+    let current_version = env!("CARGO_PKG_VERSION");
 
-        config.expect_ok_ex(
-            &["rustup", "update", "nightly"],
-            &format!(
-                r"
+    cx.config.expect_ok_ex(
+        &["rustup", "update", "nightly"],
+        &format!(
+            r"
   nightly-{} installed - 1.3.0 (hash-nightly-2)
 
 rustup - Update available : {} -> {}
 ",
-                &this_host_triple(),
-                current_version,
-                test_version
-            ),
-            for_host!(
-                r"info: syncing channel updates for 'nightly-{0}'
+            &this_host_triple(),
+            current_version,
+            test_version
+        ),
+        for_host!(
+            r"info: syncing channel updates for 'nightly-{0}'
 info: latest update on 2015-01-02, rust version 1.3.0 (hash-nightly-2)
 info: downloading component 'cargo'
 info: downloading component 'rust-docs'
@@ -68,28 +70,29 @@ info: installing component 'rust-docs'
 info: installing component 'rust-std'
 info: installing component 'rustc'
 "
-            ),
-        );
-    })
+        ),
+    );
 }
 
 #[test]
 fn update_once_and_self_update() {
     let test_version = "2.0.0";
     let mut cx = CliTestContext::from(Scenario::SimpleV2);
-    with_update_server(&mut cx.config, test_version, &|config| {
-        config.expect_ok(&["rustup-init", "-y", "--no-modify-path"]);
-        config.expect_ok(&["rustup", "set", "auto-self-update", "enable"]);
-        config.expect_ok_ex(
-            &["rustup", "update", "nightly"],
-            for_host!(
-                r"
+    let _dist_guard = cx.with_update_server(test_version);
+    cx.config
+        .expect_ok(&["rustup-init", "-y", "--no-modify-path"]);
+    cx.config
+        .expect_ok(&["rustup", "set", "auto-self-update", "enable"]);
+    cx.config.expect_ok_ex(
+        &["rustup", "update", "nightly"],
+        for_host!(
+            r"
   nightly-{0} installed - 1.3.0 (hash-nightly-2)
 
 "
-            ),
-            for_host!(
-                r"info: syncing channel updates for 'nightly-{0}'
+        ),
+        for_host!(
+            r"info: syncing channel updates for 'nightly-{0}'
 info: latest update on 2015-01-02, rust version 1.3.0 (hash-nightly-2)
 info: downloading component 'cargo'
 info: downloading component 'rust-docs'
@@ -102,9 +105,8 @@ info: installing component 'rustc'
 info: checking for self-update
 info: downloading self-update
 "
-            ),
-        );
-    });
+        ),
+    );
 }
 
 #[test]
@@ -182,32 +184,30 @@ nightly-{0} - Update available : 1.2.0 (hash-nightly-1) -> 1.3.0 (hash-nightly-2
 fn check_updates_self() {
     let test_version = "2.0.0";
     let mut cx = CliTestContext::from(Scenario::SimpleV2);
-    with_update_server(&mut cx.config, test_version, &|config| {
-        let current_version = env!("CARGO_PKG_VERSION");
+    let _dist_guard = cx.with_update_server(test_version);
+    let current_version = env!("CARGO_PKG_VERSION");
 
-        config.expect_stdout_ok(
-            &["rustup", "check"],
-            &format!(
-                r"rustup - Update available : {current_version} -> {test_version}
+    cx.config.expect_stdout_ok(
+        &["rustup", "check"],
+        &format!(
+            r"rustup - Update available : {current_version} -> {test_version}
 "
-            ),
-        );
-    });
+        ),
+    );
 }
 
 #[test]
 fn check_updates_self_no_change() {
     let current_version = env!("CARGO_PKG_VERSION");
     let mut cx = CliTestContext::from(Scenario::SimpleV2);
-    with_update_server(&mut cx.config, current_version, &|config| {
-        config.expect_stdout_ok(
-            &["rustup", "check"],
-            &format!(
-                r"rustup - Up to date : {current_version}
+    let _dist_guard = cx.with_update_server(current_version);
+    cx.config.expect_stdout_ok(
+        &["rustup", "check"],
+        &format!(
+            r"rustup - Up to date : {current_version}
 "
-            ),
-        );
-    });
+        ),
+    );
 }
 
 #[test]
