From 56c67ca4b708041c87c577d75ce9b6e1ba211347 Mon Sep 17 00:00:00 2001
From: katelyn martin <katelyn@penumbralabs.xyz>
Date: Mon, 8 Apr 2024 12:22:52 -0400
Subject: [PATCH] =?UTF-8?q?mock-consensus:=20=F0=9F=AB=9B=20`two=5Fvalidat?=
 =?UTF-8?q?ors`=20adds=20two=20keys?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

fixes #3937.

* #3937
* #3588

this adds a `two_validators` method to the test node builder, so that
tests may set up a test node that has two validator keys.
---
 crates/test/mock-consensus/src/builder.rs | 36 ++++++++++++++++++-----
 1 file changed, 29 insertions(+), 7 deletions(-)

diff --git a/crates/test/mock-consensus/src/builder.rs b/crates/test/mock-consensus/src/builder.rs
index e0fc5bae0b..80c86ca21b 100644
--- a/crates/test/mock-consensus/src/builder.rs
+++ b/crates/test/mock-consensus/src/builder.rs
@@ -8,7 +8,6 @@ mod init_chain;
 use {
     crate::{Keyring, TestNode},
     bytes::Bytes,
-    std::collections::BTreeMap,
 };
 
 /// A buider, used to prepare and instantiate a new [`TestNode`].
@@ -58,15 +57,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 }
     }
 }