Add count to move, fix tests

common/eth2/src/lighthouse_vc/std_types.rs

@@ -28,7 +28,7 @@ pub struct ListKeystoresResponse {
     pub data: Vec<SingleKeystoreResponse>,
 }
 
-#[derive(Debug, Deserialize, Serialize, PartialEq)]
+#[derive(Debug, Deserialize, Serialize, PartialEq, Eq, Hash)]
 pub struct SingleKeystoreResponse {
     pub validating_pubkey: PublicKeyBytes,
     pub derivation_path: Option<String>,

lighthouse/tests/validator_manager.rs

@@ -272,7 +272,7 @@ pub fn validator_move_misc_flags_0() {
                 src_vc_token_path: PathBuf::from("./1.json"),
                 dest_vc_url: SensitiveUrl::parse("http://localhost:2").unwrap(),
                 dest_vc_token_path: PathBuf::from("./2.json"),
-                validators: Validators::Some(vec![
+                validators: Validators::Specific(vec![
                     PublicKeyBytes::from_str(EXAMPLE_PUBKEY_0).unwrap(),
                     PublicKeyBytes::from_str(EXAMPLE_PUBKEY_1).unwrap(),
                 ]),
@@ -299,7 +299,7 @@ pub fn validator_move_misc_flags_1() {
                 src_vc_token_path: PathBuf::from("./1.json"),
                 dest_vc_url: SensitiveUrl::parse("http://localhost:2").unwrap(),
                 dest_vc_token_path: PathBuf::from("./2.json"),
-                validators: Validators::Some(vec![
+                validators: Validators::Specific(vec![
                     PublicKeyBytes::from_str(EXAMPLE_PUBKEY_0).unwrap()
                 ]),
                 builder_proposals: Some(false),
@@ -309,3 +309,26 @@ pub fn validator_move_misc_flags_1() {
             assert_eq!(expected, config);
         });
 }
+
+#[test]
+pub fn validator_move_count() {
+    CommandLineTest::validators_move()
+        .flag("--src-validator-client-url", Some("http://localhost:1"))
+        .flag("--src-validator-client-token", Some("./1.json"))
+        .flag("--dest-validator-client-url", Some("http://localhost:2"))
+        .flag("--dest-validator-client-token", Some("./2.json"))
+        .flag("--validators", Some("42"))
+        .assert_success(|config| {
+            let expected = MoveConfig {
+                src_vc_url: SensitiveUrl::parse("http://localhost:1").unwrap(),
+                src_vc_token_path: PathBuf::from("./1.json"),
+                dest_vc_url: SensitiveUrl::parse("http://localhost:2").unwrap(),
+                dest_vc_token_path: PathBuf::from("./2.json"),
+                validators: Validators::Count(42),
+                builder_proposals: None,
+                fee_recipient: None,
+                gas_limit: None,
+            };
+            assert_eq!(expected, config);
+        });
+}

validator_manager/src/validators/create_validators.rs

@@ -560,7 +560,7 @@ pub mod tests {
                 disable_deposits: false,
                 specify_voting_keystore_password: false,
                 eth1_withdrawal_address: None,
-                builder_proposals: false,
+                builder_proposals: None,
                 fee_recipient: None,
                 gas_limit: None,
                 bn_url: None,
@@ -610,8 +610,8 @@ pub mod tests {
                     assert!(validator.slashing_protection.is_none());
                     assert_eq!(validator.fee_recipient, config.fee_recipient);
                     assert_eq!(validator.gas_limit, config.gas_limit);
-                    assert_eq!(validator.builder_proposals, Some(config.builder_proposals));
+                    assert_eq!(validator.builder_proposals, config.builder_proposals);
-                    assert_eq!(validator.enabled, Some(true));
+                    assert_eq!(validator.enabled, None);
                 }
 
                 let deposits_path = output_dir.path().join(DEPOSITS_FILENAME);
@@ -759,7 +759,7 @@ pub mod tests {
         TestBuilder::default()
             .mutate_config(|config| {
                 config.deposit_gwei = 42;
-                config.builder_proposals = true;
+                config.builder_proposals = Some(true);
                 config.gas_limit = Some(1337);
             })
             .run_test()

validator_manager/src/validators/move_validators.rs

@@ -133,7 +133,8 @@ pub fn cli_app<'a, 'b>() -> App<'a, 'b> {
 #[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
 pub enum Validators {
     All,
-    Some(Vec<PublicKeyBytes>),
+    Count(usize),
+    Specific(Vec<PublicKeyBytes>),
 }
 
 impl FromStr for Validators {
@@ -142,11 +143,16 @@ impl FromStr for Validators {
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         match s {
             "all" => Ok(Validators::All),
-            other => other
+            pubkeys if pubkeys.starts_with("0x") => pubkeys
                 .split(',')
                 .map(PublicKeyBytes::from_str)
                 .collect::<Result<_, _>>()
-                .map(Validators::Some),
+                .map(Validators::Specific),
+            other => usize::from_str(other)
+                .map_err(|_| {
+                    format!("Expected \"all\", a list of 0x-prefixed pubkeys or an integer")
+                })
+                .map(Validators::Count),
         }
     }
 }
@@ -228,7 +234,26 @@ async fn run<'a>(config: MoveConfig) -> Result<(), String> {
 
     let pubkeys_to_move = match validators {
         Validators::All => src_keystores.iter().map(|v| v.validating_pubkey).collect(),
-        Validators::Some(request_pubkeys) => {
+        Validators::Count(count) => {
+            let mut viable_pubkeys: Vec<_> = src_keystores
+                .iter()
+                .filter(|v| !v.readonly.unwrap_or(true))
+                .map(|v| v.validating_pubkey)
+                .collect();
+            viable_pubkeys.sort_unstable_by_key(|pubkey| pubkey.serialize());
+            viable_pubkeys
+                .get(0..count)
+                .ok_or_else(|| {
+                    format!(
+                        "Cannot move {} keystores since source validator client only has {} \
+                        keystores which are able to be moved (not read-only).",
+                        count,
+                        viable_pubkeys.len()
+                    )
+                })?
+                .to_vec()
+        }
+        Validators::Specific(request_pubkeys) => {
             let request_pubkeys_set: HashSet<_> = request_pubkeys.iter().collect();
             let src_pubkeys_set: HashSet<_> =
                 src_keystores.iter().map(|v| &v.validating_pubkey).collect();
@@ -623,62 +648,88 @@ mod test {
                 dest_vc_url: dest_vc.url.clone(),
                 dest_vc_token_path: dest_vc_token_path.clone(),
                 validators: validators.clone(),
-                builder_proposals: false,
+                builder_proposals: None,
                 fee_recipient: None,
                 gas_limit: None,
             };
 
             let result = run(move_config).await;
 
-            let src_vc_final_keystores = src_vc_client.get_keystores().await.unwrap().data;
+            if result.is_ok() {
-            let dest_vc_final_keystores = dest_vc_client.get_keystores().await.unwrap().data;
+                let src_vc_final_keystores = src_vc_client.get_keystores().await.unwrap().data;
+                let dest_vc_final_keystores = dest_vc_client.get_keystores().await.unwrap().data;
 
-            match validators {
+                match validators {
-                Validators::All => {
+                    Validators::All => {
-                    assert!(
-                        src_vc_final_keystores.is_empty(),
-                        "all keystores should be removed from source vc"
-                    );
-                    assert_eq!(
-                        dest_vc_final_keystores.len(),
-                        dest_vc_initial_keystores.len() + src_vc_initial_keystores.len()
-                            - self.duplicates,
-                        "the correct count of keystores should have been moved to the dest"
-                    );
-                    for initial_keystore in &src_vc_initial_keystores {
                         assert!(
-                            dest_vc_final_keystores.contains(initial_keystore),
+                            src_vc_final_keystores.is_empty(),
-                            "the source keystore should be present at the dest"
+                            "all keystores should be removed from source vc"
-                        )
+                        );
+                        assert_eq!(
+                            dest_vc_final_keystores.len(),
+                            dest_vc_initial_keystores.len() + src_vc_initial_keystores.len()
+                                - self.duplicates,
+                            "the correct count of keystores should have been moved to the dest"
+                        );
+                        for initial_keystore in &src_vc_initial_keystores {
+                            assert!(
+                                dest_vc_final_keystores.contains(initial_keystore),
+                                "the source keystore should be present at the dest"
+                            )
+                        }
                     }
-                }
+                    Validators::Count(count) => {
-                Validators::Some(pubkeys) => {
+                        assert_eq!(
-                    assert_eq!(
+                            src_vc_final_keystores.len(),
-                        src_vc_final_keystores.len(),
+                            src_vc_initial_keystores.len() - count,
-                        src_vc_initial_keystores
+                            "keystores should be removed from source vc"
-                            .len()
-                            .checked_sub(pubkeys.len())
-                            .unwrap(),
-                        "the correct count of validators should have been removed from the src"
-                    );
-                    assert_eq!(
-                        dest_vc_final_keystores.len(),
-                        dest_vc_initial_keystores.len() + pubkeys.len() - self.duplicates,
-                        "the correct count of keystores should have been moved to the dest"
-                    );
-                    for pubkey in pubkeys {
-                        let initial_keystore = src_vc_initial_keystores
-                            .iter()
-                            .find(|k| k.validating_pubkey == pubkey)
-                            .unwrap();
-                        assert!(
-                            !src_vc_final_keystores.contains(initial_keystore),
-                            "the keystore should not be present at the source"
                         );
-                        assert!(
+                        assert_eq!(
-                            dest_vc_final_keystores.contains(initial_keystore),
+                            dest_vc_final_keystores.len(),
-                            "the keystore should be present at the dest"
+                            dest_vc_initial_keystores.len() + count - self.duplicates,
+                            "the correct count of keystores should have been moved to the dest"
                         );
+                        let moved_keystores: Vec<_> = {
+                            let initial_set: HashSet<_> = src_vc_initial_keystores.iter().collect();
+                            let final_set: HashSet<_> = src_vc_final_keystores.iter().collect();
+                            initial_set.difference(&final_set).cloned().collect()
+                        };
+                        assert_eq!(moved_keystores.len(), count);
+                        for moved_keystore in &moved_keystores {
+                            assert!(
+                                dest_vc_final_keystores.contains(moved_keystore),
+                                "the moved keystore should be present at the dest"
+                            )
+                        }
+                    }
+                    Validators::Specific(pubkeys) => {
+                        assert_eq!(
+                            src_vc_final_keystores.len(),
+                            src_vc_initial_keystores
+                                .len()
+                                .checked_sub(pubkeys.len())
+                                .unwrap(),
+                            "the correct count of validators should have been removed from the src"
+                        );
+                        assert_eq!(
+                            dest_vc_final_keystores.len(),
+                            dest_vc_initial_keystores.len() + pubkeys.len() - self.duplicates,
+                            "the correct count of keystores should have been moved to the dest"
+                        );
+                        for pubkey in pubkeys {
+                            let initial_keystore = src_vc_initial_keystores
+                                .iter()
+                                .find(|k| k.validating_pubkey == pubkey)
+                                .unwrap();
+                            assert!(
+                                !src_vc_final_keystores.contains(initial_keystore),
+                                "the keystore should not be present at the source"
+                            );
+                            assert!(
+                                dest_vc_final_keystores.contains(initial_keystore),
+                                "the keystore should be present at the dest"
+                            );
+                        }
                     }
                 }
             }
@@ -697,6 +748,10 @@ mod test {
             assert_eq!(self.result, Ok(()))
         }
 
+        fn assert_err(self) {
+            assert!(self.result.is_err())
+        }
+
         fn assert_err_is(self, msg: String) {
             assert_eq!(self.result, Err(msg))
         }
@@ -728,7 +783,7 @@ mod test {
             .await
             .with_src_validators(1, 0)
             .await
-            .run_test(|pubkeys| Validators::Some(pubkeys.to_vec()))
+            .run_test(|pubkeys| Validators::Specific(pubkeys.to_vec()))
             .await
             .assert_ok();
     }
@@ -769,7 +824,7 @@ mod test {
             .with_dest_validators(2, 0)
             .await
             .register_duplicates(1)
-            .run_test(|pubkeys| Validators::Some(pubkeys[0..1].to_vec()))
+            .run_test(|pubkeys| Validators::Specific(pubkeys[0..1].to_vec()))
             .await
             .assert_ok();
     }
@@ -791,7 +846,7 @@ mod test {
             .await
             .with_src_validators(3, 0)
             .await
-            .run_test(|pubkeys| Validators::Some(pubkeys[0..1].to_vec()))
+            .run_test(|pubkeys| Validators::Specific(pubkeys[0..1].to_vec()))
             .await
             .assert_ok();
     }
@@ -802,7 +857,7 @@ mod test {
             .await
             .with_src_validators(3, 0)
             .await
-            .run_test(|pubkeys| Validators::Some(pubkeys[0..2].to_vec()))
+            .run_test(|pubkeys| Validators::Specific(pubkeys[0..2].to_vec()))
             .await
             .assert_ok();
     }
@@ -813,8 +868,41 @@ mod test {
             .await
             .with_src_validators(3, 42)
             .await
-            .run_test(|pubkeys| Validators::Some(pubkeys.to_vec()))
+            .run_test(|pubkeys| Validators::Specific(pubkeys.to_vec()))
             .await
             .assert_ok();
     }
+
+    #[tokio::test]
+    async fn three_validators_move_one_by_count() {
+        TestBuilder::new()
+            .await
+            .with_src_validators(3, 0)
+            .await
+            .run_test(|_| Validators::Count(1))
+            .await
+            .assert_ok();
+    }
+
+    #[tokio::test]
+    async fn three_validators_move_two_by_count() {
+        TestBuilder::new()
+            .await
+            .with_src_validators(3, 0)
+            .await
+            .run_test(|_| Validators::Count(2))
+            .await
+            .assert_ok();
+    }
+
+    #[tokio::test]
+    async fn one_validators_move_two_by_count() {
+        TestBuilder::new()
+            .await
+            .with_src_validators(1, 0)
+            .await
+            .run_test(|_| Validators::Count(2))
+            .await
+            .assert_err();
+    }
 }