From b58c6c9840f4060785dcc5a917a4c77513e4b6bd Mon Sep 17 00:00:00 2001
From: rikosellic <64517311+rikosellic@users.noreply.github.com>
Date: Mon, 5 Jan 2026 14:12:39 +0800
Subject: [PATCH] Replace test_branch with branch

---
 src/main.rs  | 17 +++++++----------
 src/verus.rs |  7 ++++++-
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 2e566dc..6257ba4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -108,11 +108,12 @@ struct BootstrapArgs {
     )]
     debug: bool,
 
-    #[arg(long = "test_branch", 
-        help = "Use the test branch of our Verus fork for testing",
-        default_value = "false", 
-        action = ArgAction::SetTrue)]
-    test_branch: bool,
+    #[arg(
+        long = "branch",
+        help = "The branch name to pull (only used with --upgrade)",
+        value_name = "BRANCH_NAME"
+    )]
+    branch: Option<String>,
 }
 
 #[derive(Parser, Debug)]
@@ -424,11 +425,7 @@ fn bootstrap(args: &BootstrapArgs) -> Result<(), DynError> {
     let options = verus::install::VerusInstallOpts {
         release: !args.debug,
         restart: args.restart,
-        branch: if args.test_branch {
-            Some("update-test".into())
-        } else {
-            None
-        },
+        branch: args.branch.clone(),
         force_reset: args.upgrade,
     };
 
diff --git a/src/verus.rs b/src/verus.rs
index d8cd2a7..318dc52 100644
--- a/src/verus.rs
+++ b/src/verus.rs
@@ -1562,7 +1562,12 @@ pub mod install {
 
         // Find the matching remote branch
         let upstream_branch = format!("refs/remotes/origin/{}", target_branch);
-        let upstream_ref = repo.find_reference(&upstream_branch)?;
+        let upstream_ref = repo.find_reference(&upstream_branch).map_err(|_| {
+            format!(
+                "Branch '{}' does not exist in the remote repository. Please check the branch name.",
+                target_branch
+            )
+        })?;
         let upstream_commit = upstream_ref.peel_to_commit()?;
 
         // Check merge analysis