Add --count flag as per @michaelsproul's comment

paulhauner · Jul 11, 2023 · fbb844d · fbb844d
1 parent 4af863a
commit fbb844d
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 25 deletions.
diff --git a/lighthouse/tests/validator_manager.rs b/lighthouse/tests/validator_manager.rs
@@ -321,7 +321,7 @@ pub fn validator_move_count() {
         .flag("--src-vc-token", Some("./1.json"))
         .flag("--dest-vc-url", Some("http://localhost:2"))
         .flag("--dest-vc-token", Some("./2.json"))
-        .flag("--validators", Some("42"))
+        .flag("--count", Some("42"))
         .assert_success(|config| {
             let expected = MoveConfig {
                 src_vc_url: SensitiveUrl::parse("http://localhost:1").unwrap(),

diff --git a/validator_client/src/http_api/tests.rs b/validator_client/src/http_api/tests.rs
@@ -33,7 +33,6 @@ use std::sync::Arc;
 use std::time::Duration;
 use task_executor::test_utils::TestRuntime;
 use tempfile::{tempdir, TempDir};
-use tokio::runtime::Runtime;
 use types::graffiti::GraffitiString;
 
 const PASSWORD_BYTES: &[u8] = &[42, 50, 37];

diff --git a/validator_manager/src/common.rs b/validator_manager/src/common.rs
@@ -16,6 +16,7 @@ use types::*;
 
 pub const IGNORE_DUPLICATES_FLAG: &str = "ignore-duplicates";
 pub const STDIN_INPUTS_FLAG: &str = "stdin-inputs";
+pub const COUNT_FLAG: &str = "count";
 
 /// When the `ethereum/staking-deposit-cli` tool generates deposit data JSON, it adds a
 /// `deposit_cli_version` to protect the web-based "Launchpad" tool against a breaking change that

diff --git a/validator_manager/src/create_validators.rs b/validator_manager/src/create_validators.rs
@@ -18,7 +18,6 @@ pub const CMD: &str = "create";
 pub const OUTPUT_PATH_FLAG: &str = "output-path";
 pub const DEPOSIT_GWEI_FLAG: &str = "deposit-gwei";
 pub const DISABLE_DEPOSITS_FLAG: &str = "disable-deposits";
-pub const COUNT_FLAG: &str = "count";
 pub const FIRST_INDEX_FLAG: &str = "first-index";
 pub const MNEMONIC_FLAG: &str = "mnemonic-path";
 pub const SPECIFY_VOTING_KEYSTORE_PASSWORD_FLAG: &str = "specify-voting-keystore-password";

diff --git a/validator_manager/src/move_validators.rs b/validator_manager/src/move_validators.rs
@@ -120,7 +120,14 @@ pub fn cli_app<'a, 'b>() -> App<'a, 'b> {
                     validator pubkeys, an integer count of validators or the \
                     keyword \"all\".",
                 )
-                .required(true)
+                .takes_value(true),
+        )
+        .arg(
+            Arg::with_name(COUNT_FLAG)
+                .long(COUNT_FLAG)
+                .value_name("VALIDATOR_COUNT")
+                .help("The number of validators to move.")
+                .conflicts_with(VALIDATORS_FLAG)
                 .takes_value(true),
         )
         .arg(
@@ -172,26 +179,6 @@ pub enum Validators {
     Specific(Vec<PublicKeyBytes>),
 }
 
-impl FromStr for Validators {
-    type Err = String;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "all" => Ok(Validators::All),
-            pubkeys if pubkeys.starts_with("0x") => pubkeys
-                .split(',')
-                .map(PublicKeyBytes::from_str)
-                .collect::<Result<_, _>>()
-                .map(Validators::Specific),
-            other => usize::from_str(other)
-                .map_err(|_| {
-                    "Expected \"all\", a list of 0x-prefixed pubkeys or an integer".to_string()
-                })
-                .map(Validators::Count),
-        }
-    }
-}
-
 #[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
 pub struct MoveConfig {
     pub src_vc_url: SensitiveUrl,
@@ -207,12 +194,30 @@ pub struct MoveConfig {
 
 impl MoveConfig {
     fn from_cli(matches: &ArgMatches) -> Result<Self, String> {
+        let count_flag = clap_utils::parse_optional(matches, COUNT_FLAG)?;
+        let validators_flag = matches.value_of(VALIDATORS_FLAG);
+        let validators = match (count_flag, validators_flag) {
+            (Some(count), None) => Validators::Count(count),
+            (None, Some(string)) => match string {
+                "all" => Validators::All,
+                pubkeys => pubkeys
+                    .split(',')
+                    .map(PublicKeyBytes::from_str)
+                    .collect::<Result<Vec<_>, _>>()
+                    .map(Validators::Specific)?,
+            },
+            (None, None) => Err("Must supply either --{:VALIDATORS_FLAG} or --{:COUNT_FLAG}.")?,
+            (Some(_), Some(_)) => {
+                Err("Cannot supply both --{:VALIDATORS_FLAG} and --{:COUNT_FLAG}.")?
+            }
+        };
+
         Ok(Self {
             src_vc_url: clap_utils::parse_required(matches, SRC_VC_URL_FLAG)?,
             src_vc_token_path: clap_utils::parse_required(matches, SRC_VC_TOKEN_FLAG)?,
             dest_vc_url: clap_utils::parse_required(matches, DEST_VC_URL_FLAG)?,
             dest_vc_token_path: clap_utils::parse_required(matches, DEST_VC_TOKEN_FLAG)?,
-            validators: clap_utils::parse_required(matches, VALIDATORS_FLAG)?,
+            validators,
             builder_proposals: clap_utils::parse_optional(matches, BUILDER_PROPOSALS_FLAG)?,
             fee_recipient: clap_utils::parse_optional(matches, FEE_RECIPIENT_FLAG)?,
             gas_limit: clap_utils::parse_optional(matches, GAS_LIMIT_FLAG)?,