diff --git a/crates/test/mock-consensus/src/builder.rs b/crates/test/mock-consensus/src/builder.rs index e0fc5bae0b..09c9a9a894 100644 --- a/crates/test/mock-consensus/src/builder.rs +++ b/crates/test/mock-consensus/src/builder.rs @@ -58,15 +58,38 @@ impl Builder { ); } - // Generate a consensus key. + // Generate a key and place it in the keyring. + let mut keyring = Keyring::new(); + Self::add_key(&mut keyring); + + Self { keyring, ..self } + } + + /// Generates a pair of validator keys. + pub fn two_validators(self) -> Self { + let Self { keyring: prev, .. } = self; + + // Log a warning if we are about to overwrite any existing keys. + if !prev.is_empty() { + tracing::warn!( + count = %prev.len(), + "builder overwriting entries in keyring, this may be a bug!" + ); + } + + // Generate two keys and place them in the keyring. + let mut keyring = Keyring::new(); + Self::add_key(&mut keyring); + Self::add_key(&mut keyring); + + Self { keyring, ..self } + } + + /// Generates consensus keys and places them in the provided keyring. + fn add_key(keyring: &mut Keyring) { let sk = ed25519_consensus::SigningKey::new(rand_core::OsRng); let vk = sk.verification_key(); tracing::trace!(verification_key = ?vk, "generated consensus key"); - - // Place it into the keyring. - let mut keyring = BTreeMap::new(); keyring.insert(vk, sk); - - Self { keyring, ..self } } }