mock-consensus: 🔙 TestNode<C> has an on_block callback

this commit introduces a new, optional component of the mock consensus
test node. `on_block` is a `tendermint::Block -> ()` function that can
be used to collect the blocks being generated by the test node.

crates/test/mock-consensus/src/block.rs

@@ -113,7 +113,7 @@ where
             evidence: _,
             last_commit,
             ..
-        } = block.tap(|block| {
+        } = block.clone().tap(|block| {
             tracing::span::Span::current()
                 .record("height", block.header.height.value())
                 .record("time", block.header.time.unix_timestamp());
@@ -130,6 +130,9 @@ where
         test_node.commit().await?;
         trace!("finished sending block");
 
+        // If an `on_block` callback was set, call it now.
+        test_node.on_block.as_mut().map(move |f| f(block));
+
         Ok(())
     }
 

crates/test/mock-consensus/src/builder.rs

@@ -6,7 +6,7 @@
 mod init_chain;
 
 use {
-    crate::{Keyring, TestNode},
+    crate::{Keyring, OnBlockFn, TestNode},
     bytes::Bytes,
 };
 
@@ -15,6 +15,7 @@ use {
 pub struct Builder {
     pub app_state: Option<Bytes>,
     pub keyring: Keyring,
+    pub on_block: Option<OnBlockFn>,
 }
 
 impl TestNode<()> {
@@ -91,4 +92,15 @@ impl Builder {
         tracing::trace!(verification_key = ?vk, "generated consensus key");
         keyring.insert(vk, sk);
     }
+
+    /// Sets a callback that will be invoked when a new block is constructed.
+    pub fn on_block<F>(self, f: F) -> Self
+    where
+        F: FnMut(tendermint::Block) + Send + Sync + 'static,
+    {
+        Self {
+            on_block: Some(Box::new(f)),
+            ..self
+        }
+    }
 }

crates/test/mock-consensus/src/builder/init_chain.rs

@@ -38,6 +38,7 @@ impl Builder {
         let Self {
             app_state: Some(app_state),
             keyring,
+            on_block,
         } = self
         else {
             bail!("builder was not fully initialized")
@@ -69,6 +70,7 @@ impl Builder {
             height: block::Height::from(0_u8),
             last_app_hash: app_hash.as_bytes().to_owned(),
             keyring,
+            on_block,
         })
     }
 

crates/test/mock-consensus/src/lib.rs

@@ -82,8 +82,13 @@ pub struct TestNode<C> {
     ///
     /// Entries in this keyring consist of a [`VerificationKey`] and a [`SigningKey`].
     keyring: Keyring,
+    /// A callback that will be invoked when a new block is constructed.
+    on_block: Option<OnBlockFn>,
 }
 
+/// A type alias for the `TestNode::on_block` callback.
+pub type OnBlockFn = Box<dyn FnMut(tendermint::Block) + Send + Sync + 'static>;
+
 /// An ordered map of consensus keys.
 ///
 /// Entries in this keyring consist of a [`VerificationKey`] and a [`SigningKey`].