Add ability to set fixed time and block duration callback in test node

crates/core/app/tests/app_blocktimes_increment.rs

@@ -0,0 +1,99 @@
+use {
+    self::common::BuilderExt,
+    cnidarium::TempStorage,
+    penumbra_app::{
+        genesis::{self, AppState},
+        server::consensus::Consensus,
+    },
+    penumbra_mock_consensus::TestNode,
+    std::time::Duration,
+    tap::{Tap, TapFallible},
+};
+
+mod common;
+
+/// This is more of a test of test code; this ensures
+/// that the mock tendermint's block times are set and increment as expected.
+#[tokio::test]
+async fn mock_tendermint_block_times_correct() -> anyhow::Result<()> {
+    // Install a test logger, and acquire some temporary storage.
+    let guard = common::set_tracing_subscriber();
+    let storage = TempStorage::new().await?;
+
+    // Fixed start time:
+    let start_time = tendermint::Time::parse_from_rfc3339("2022-02-11T17:30:50.425417198Z")?;
+
+    // Define our application state, and start the test node.
+    let mut test_node = {
+        let app_state = AppState::Content(
+            genesis::Content::default().with_chain_id(TestNode::<()>::CHAIN_ID.to_string()),
+        );
+        let consensus = Consensus::new(storage.as_ref().clone());
+        // This should use the default time callback of 5s
+        TestNode::builder()
+            .single_validator()
+            .with_penumbra_auto_app_state(app_state)?
+            .with_initial_timestamp(start_time)
+            .init_chain(consensus)
+            .await
+            .tap_ok(|e| tracing::info!(hash = %e.last_app_hash_hex(), "finished init chain"))?
+    };
+
+    // The test node's time should be the initial timestamp before any blocks are committed
+    assert_eq!(*test_node.timestamp(), start_time);
+
+    // Test a handful of block executions
+    for i in 0..10 {
+        // Execute a block on the test node
+        test_node.block().execute().await?;
+
+        // Ensure the time has incremented by 5 seconds
+        assert_eq!(
+            *test_node.timestamp(),
+            start_time
+                .checked_add(Duration::from_secs(5 * (i + 1)))
+                .unwrap()
+        );
+    }
+
+    // Now do it with a different duration.
+    let block_duration = Duration::from_secs(13);
+    let storage = TempStorage::new().await?;
+    let mut test_node = {
+        let app_state = AppState::Content(
+            genesis::Content::default().with_chain_id(TestNode::<()>::CHAIN_ID.to_string()),
+        );
+        let consensus = Consensus::new(storage.as_ref().clone());
+        // This should use the default time callback of 5s
+        TestNode::builder()
+            .single_validator()
+            .with_penumbra_auto_app_state(app_state)?
+            .with_initial_timestamp(start_time)
+            // Set a callback to add 13 seconds instead
+            .ts_callback(move |t| t.checked_add(block_duration).unwrap())
+            .init_chain(consensus)
+            .await
+            .tap_ok(|e| tracing::info!(hash = %e.last_app_hash_hex(), "finished init chain"))?
+    };
+
+    // The test node's time should be the initial timestamp before any blocks are committed
+    assert_eq!(*test_node.timestamp(), start_time);
+
+    // Test a handful of block executions
+    for i in 0..10 {
+        // Execute a block on the test node
+        test_node.block().execute().await?;
+
+        // Ensure the time has incremented by 5 seconds
+        assert_eq!(
+            *test_node.timestamp(),
+            start_time.checked_add(block_duration * (i + 1)).unwrap()
+        );
+    }
+
+    // Free our temporary storage.
+    Ok(())
+        .tap(|_| drop(test_node))
+        .tap(|_| drop(storage))
+        .tap(|_| drop(guard))
+}

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

@@ -10,7 +10,7 @@ use {
         block::{self, header::Version, Block, Commit, Header, Round},
         chain, evidence,
         v0_37::abci::{ConsensusRequest, ConsensusResponse},
-        AppHash, Hash,
+        AppHash, Hash, Time,
     },
     tower::{BoxError, Service},
     tracing::{instrument, trace},
@@ -36,6 +36,8 @@ pub struct Builder<'e, C> {
     evidence: evidence::List,
     /// The list of signatures.
     signatures: Vec<block::CommitSig>,
+    /// The timestamp of the block.
+    timestamp: Time,
 }
 
 // === impl TestNode ===
@@ -47,12 +49,15 @@ impl<C> TestNode<C> {
     /// included in the block. Use [`Builder::with_signatures()`] to set a different set of
     /// validator signatures.
     pub fn block(&mut self) -> Builder<'_, C> {
+        let ts = self.timestamp.clone();
         let signatures = self.generate_signatures().collect();
+        // set default TS hook
         Builder {
             test_node: self,
             data: Default::default(),
             evidence: Default::default(),
             signatures,
+            timestamp: ts,
         }
     }
 }
@@ -133,6 +138,9 @@ where
         // If an `on_block` callback was set, call it now.
         test_node.on_block.as_mut().map(move |f| f(block));
 
+        // Call the timestamp callback to increment the node's current timestamp.
+        test_node.timestamp = (test_node.ts_callback)(test_node.timestamp.clone());
+
         Ok(())
     }
 
@@ -149,6 +157,7 @@ where
             evidence,
             test_node,
             signatures,
+            timestamp,
         } = self;
 
         let height = {
@@ -177,7 +186,7 @@ where
             version: Version { block: 1, app: 1 },
             chain_id: chain::Id::try_from("test".to_owned())?,
             height,
-            time: tendermint::Time::now(),
+            time: timestamp,
             last_block_id: None,
             last_commit_hash: None,
             data_hash: None,

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

@@ -17,10 +17,10 @@ mod sign {
     /// Returns a [commit signature] saying this validator voted for the block.
     ///
     /// [commit signature]: CommitSig
-    pub(super) fn commit(validator_address: Id) -> CommitSig {
+    pub(super) fn commit(validator_address: Id, timestamp: Time) -> CommitSig {
         CommitSig::BlockIdFlagCommit {
             validator_address,
-            timestamp: timestamp(),
+            timestamp,
             signature: None,
         }
     }
@@ -29,21 +29,13 @@ mod sign {
     ///
     /// [commit signature]: CommitSig
     #[allow(dead_code)]
-    pub(super) fn nil(validator_address: Id) -> CommitSig {
+    pub(super) fn nil(validator_address: Id, timestamp: Time) -> CommitSig {
         CommitSig::BlockIdFlagNil {
             validator_address,
-            timestamp: timestamp(),
+            timestamp,
             signature: None,
         }
     }
-
-    /// Generates a new timestamp, marked at the current time.
-    //
-    //  TODO(kate): see https://github.com/penumbra-zone/penumbra/issues/3759, re: timestamps.
-    //              eventually, we will add hooks so that we can control these timestamps.
-    fn timestamp() -> Time {
-        Time::now()
-    }
 }
 
 // === impl TestNode ===
@@ -58,12 +50,14 @@ impl<C> TestNode<C> {
         self.keyring
             .keys()
             .map(|vk| {
-                <Sha256 as Digest>::digest(vk).as_slice()[0..20]
-                    .try_into()
-                    .expect("")
+                (
+                    <Sha256 as Digest>::digest(vk).as_slice()[0..20]
+                        .try_into()
+                        .expect(""),
+                    self.timestamp.clone(),
+                )
             })
-            .map(account::Id::new)
-            .map(self::sign::commit)
+            .map(|(a, b)| (self::sign::commit(account::Id::new(a), b)))
     }
 }
 

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

@@ -6,16 +6,26 @@
 mod init_chain;
 
 use {
-    crate::{Keyring, OnBlockFn, TestNode},
+    crate::{Keyring, OnBlockFn, TestNode, TsCallbackFn},
     bytes::Bytes,
+    std::time::Duration,
+    tendermint::Time,
 };
 
+// Default timestamp callback will increment the time by 5 seconds.
+// can't be const :(
+fn default_ts_callback(t: Time) -> Time {
+    t.checked_add(Duration::from_secs(5)).unwrap()
+}
+
 /// A builder, used to prepare and instantiate a new [`TestNode`].
 #[derive(Default)]
 pub struct Builder {
     pub app_state: Option<Bytes>,
     pub keyring: Keyring,
     pub on_block: Option<OnBlockFn>,
+    pub ts_callback: Option<TsCallbackFn>,
+    pub initial_timestamp: Option<Time>,
 }
 
 impl TestNode<()> {
@@ -103,4 +113,25 @@ impl Builder {
             ..self
         }
     }
+
+    /// Sets a callback that will be invoked when a block is committed, to increment
+    /// the timestamp.
+    pub fn ts_callback<F>(self, f: F) -> Self
+    where
+        F: Fn(Time) -> Time + Send + Sync + 'static,
+    {
+        Self {
+            ts_callback: Some(Box::new(f)),
+            ..self
+        }
+    }
+
+    /// Sets the starting time for the test node. If not called,
+    /// the current timestamp will be used.
+    pub fn with_initial_timestamp(self, initial_time: Time) -> Self {
+        Self {
+            initial_timestamp: Some(initial_time),
+            ..self
+        }
+    }
 }

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

@@ -39,6 +39,8 @@ impl Builder {
             app_state: Some(app_state),
             keyring,
             on_block,
+            initial_timestamp,
+            ts_callback,
         } = self
         else {
             bail!("builder was not fully initialized")
@@ -71,6 +73,8 @@ impl Builder {
             last_app_hash: app_hash.as_bytes().to_owned(),
             keyring,
             on_block,
+            timestamp: initial_timestamp.unwrap_or(Time::now()),
+            ts_callback: ts_callback.unwrap_or(Box::new(default_ts_callback)),
         })
     }
 

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

@@ -32,6 +32,7 @@
 use {
     ed25519_consensus::{SigningKey, VerificationKey},
     std::collections::BTreeMap,
+    tendermint::Time,
 };
 
 pub mod block;
@@ -84,11 +85,18 @@ pub struct TestNode<C> {
     keyring: Keyring,
     /// A callback that will be invoked when a new block is constructed.
     on_block: Option<OnBlockFn>,
+    /// A callback that will be invoked when a new block is committed, to produce the next timestamp.
+    ts_callback: TsCallbackFn,
+    /// The current timestamp of the node.
+    timestamp: Time,
 }
 
 /// A type alias for the `TestNode::on_block` callback.
 pub type OnBlockFn = Box<dyn FnMut(tendermint::Block) + Send + Sync + 'static>;
 
+/// A type alias for the `TestNode::ts_callback` callback.
+pub type TsCallbackFn = Box<dyn Fn(Time) -> Time + Send + Sync + 'static>;
+
 /// An ordered map of consensus keys.
 ///
 /// Entries in this keyring consist of a [`VerificationKey`] and a [`SigningKey`].
@@ -104,6 +112,11 @@ impl<C> TestNode<C> {
         &self.last_app_hash
     }
 
+    /// Returns the last `timestamp` value.
+    pub fn timestamp(&self) -> &Time {
+        &self.timestamp
+    }
+
     /// Returns the last `app_hash` value, represented as a hexadecimal string.
     pub fn last_app_hash_hex(&self) -> String {
         // Use upper-case hexadecimal integers, include leading zeroes.

crates/test/mock-tendermint-proxy/Cargo.toml

@@ -8,9 +8,11 @@ homepage.workspace   = true
 license.workspace    = true
 
 [dependencies]
+pbjson-types            = { workspace = true }
 penumbra-mock-consensus = { workspace = true }
 penumbra-proto          = { workspace = true, features = ["rpc", "tendermint"] }
 tap                     = { workspace = true }
 tendermint              = { workspace = true }
+tendermint-proto        = { workspace = true }
 tonic                   = { workspace = true }
 tracing                 = { workspace = true }