From a2fbe92a252955e292e36e30580b64d9fd6f8a9f Mon Sep 17 00:00:00 2001
From: SteveLauC <stevelauc@outlook.com>
Date: Wed, 18 Oct 2023 12:19:53 +0800
Subject: [PATCH] refactor: make SelfUpdate a step (#585)

---
 src/config.rs      |  1 +
 src/main.rs        | 21 ++++-------
 src/self_update.rs | 86 ++++++++++++++++++++++++++--------------------
 3 files changed, 56 insertions(+), 52 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 10ebbbb9..65e7aaa3 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -127,6 +127,7 @@ pub enum Step {
     Rustup,
     Scoop,
     Sdkman,
+    SelfUpdate,
     Sheldon,
     Shell,
     Snap,
diff --git a/src/main.rs b/src/main.rs
index c9c5f391..cc79d01d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -108,7 +108,7 @@ fn run() -> Result<()> {
     debug!("OS: {}", env!("TARGET"));
     debug!("{:?}", std::env::args());
     debug!("Binary path: {:?}", std::env::current_exe());
-    debug!("Self Update: {:?}", cfg!(feature = "self-update"));
+    debug!("self-update Feature Enabled: {:?}", cfg!(feature = "self-update"));
     debug!("Configuration: {:?}", config);
 
     if config.run_in_tmux() && env::var("TOPGRADE_INSIDE_TMUX").is_err() {
@@ -132,22 +132,15 @@ fn run() -> Result<()> {
     let ctx = execution_context::ExecutionContext::new(run_type, sudo, &git, &config);
     let mut runner = runner::Runner::new(&ctx);
 
+    // Self-Update step, this will execute only if:
+    // 1. the `self-update` feature is enabled
+    // 2. it is not disabled from configuration (env var/CLI opt/file)
     #[cfg(feature = "self-update")]
     {
-        let config_self_upgrade = env::var("TOPGRADE_NO_SELF_UPGRADE").is_err() && !config.no_self_update();
+        let should_self_update = env::var("TOPGRADE_NO_SELF_UPGRADE").is_err() && !config.no_self_update();
 
-        if !run_type.dry() && config_self_upgrade {
-            let result = self_update::self_update();
-
-            if let Err(e) = &result {
-                #[cfg(windows)]
-                {
-                    if e.downcast_ref::<Upgraded>().is_some() {
-                        return result;
-                    }
-                }
-                print_warning(format!("Self update error: {e}"));
-            }
+        if should_self_update {
+            runner.execute(Step::SelfUpdate, "Self Update", || self_update::self_update(&ctx))?;
         }
     }
 
diff --git a/src/self_update.rs b/src/self_update.rs
index fe439399..d1d664d8 100644
--- a/src/self_update.rs
+++ b/src/self_update.rs
@@ -3,6 +3,7 @@ use std::env;
 use std::os::unix::process::CommandExt as _;
 use std::process::Command;
 
+use crate::config::Step;
 use color_eyre::eyre::{bail, Result};
 use self_update_crate::backends::github::Update;
 use self_update_crate::update::UpdateStatus;
@@ -11,52 +12,61 @@ use super::terminal::*;
 #[cfg(windows)]
 use crate::error::Upgraded;
 
-pub fn self_update() -> Result<()> {
+use crate::execution_context::ExecutionContext;
+
+pub fn self_update(ctx: &ExecutionContext) -> Result<()> {
     print_separator("Self update");
-    let current_exe = env::current_exe();
 
-    let target = self_update_crate::get_target();
-    let result = Update::configure()
-        .repo_owner("topgrade-rs")
-        .repo_name("topgrade")
-        .target(target)
-        .bin_name(if cfg!(windows) { "topgrade.exe" } else { "topgrade" })
-        .show_output(false)
-        .show_download_progress(true)
-        .current_version(self_update_crate::cargo_crate_version!())
-        .no_confirm(true)
-        .build()?
-        .update_extended()?;
-
-    if let UpdateStatus::Updated(release) = &result {
-        println!("\nTopgrade upgraded to {}:\n", release.version);
-        if let Some(body) = &release.body {
-            println!("{body}");
-        }
+    if ctx.run_type().dry() {
+        println!("Would self-update");
+        Ok(())
     } else {
-        println!("Topgrade is up-to-date");
-    }
+        let assume_yes = ctx.config().yes(Step::SelfUpdate);
+        let current_exe = env::current_exe();
 
-    {
-        if result.updated() {
-            print_warning("Respawning...");
-            let mut command = Command::new(current_exe?);
-            command.args(env::args().skip(1)).env("TOPGRADE_NO_SELF_UPGRADE", "");
+        let target = self_update_crate::get_target();
+        let result = Update::configure()
+            .repo_owner("topgrade-rs")
+            .repo_name("topgrade")
+            .target(target)
+            .bin_name(if cfg!(windows) { "topgrade.exe" } else { "topgrade" })
+            .show_output(true)
+            .show_download_progress(true)
+            .current_version(self_update_crate::cargo_crate_version!())
+            .no_confirm(assume_yes)
+            .build()?
+            .update_extended()?;
 
-            #[cfg(unix)]
-            {
-                let err = command.exec();
-                bail!(err);
+        if let UpdateStatus::Updated(release) = &result {
+            println!("\nTopgrade upgraded to {}:\n", release.version);
+            if let Some(body) = &release.body {
+                println!("{body}");
             }
+        } else {
+            println!("Topgrade is up-to-date");
+        }
 
-            #[cfg(windows)]
-            {
-                #[allow(clippy::disallowed_methods)]
-                let status = command.status()?;
-                bail!(Upgraded(status));
+        {
+            if result.updated() {
+                print_info("Respawning...");
+                let mut command = Command::new(current_exe?);
+                command.args(env::args().skip(1)).env("TOPGRADE_NO_SELF_UPGRADE", "");
+
+                #[cfg(unix)]
+                {
+                    let err = command.exec();
+                    bail!(err);
+                }
+
+                #[cfg(windows)]
+                {
+                    #[allow(clippy::disallowed_methods)]
+                    let status = command.status()?;
+                    bail!(Upgraded(status));
+                }
             }
         }
-    }
 
-    Ok(())
+        Ok(())
+    }
 }