penumbra_sdk_mock_consensus/
builder.rs

1//! [`Builder`] interfaces, for creating new [`TestNode`]s.
2
3/// [`Builder`] interfaces for chain initialization.
4///
5/// Most importantly, defines [`Builder::init_chain()`].
6mod init_chain;
7
8use {
9    crate::{Keyring, OnBlockFn, TestNode, TsCallbackFn},
10    anyhow::Result,
11    bytes::Bytes,
12    ed25519_consensus::{SigningKey, VerificationKey},
13    std::time::Duration,
14    tendermint::{Genesis, Time},
15};
16
17// Default timestamp callback will increment the time by 5 seconds.
18// can't be const :(
19fn default_ts_callback(t: Time) -> Time {
20    t.checked_add(Duration::from_secs(5)).unwrap()
21}
22
23/// A builder, used to prepare and instantiate a new [`TestNode`].
24#[derive(Default)]
25pub struct Builder {
26    pub app_state: Option<Bytes>,
27    pub keyring: Keyring,
28    pub on_block: Option<OnBlockFn>,
29    pub ts_callback: Option<TsCallbackFn>,
30    pub initial_timestamp: Option<Time>,
31    pub chain_id: Option<String>,
32    /// Hardcodes a genesis to be used for the chain.
33    /// Useful if you're trying to test cometbft compatibility
34    /// but should not be used typically.
35    pub hardcoded_genesis: Option<Genesis>,
36    /// Specifies hardcoded keys to be assigned to validators.
37    /// Lowest indices are assigned first. If not enough are provided,
38    /// random keys will be generated for the remaining validators.
39    /// The default behavior is to generate random keys if none are
40    /// supplied.
41    pub keys: Vec<(SigningKey, VerificationKey)>,
42}
43
44impl TestNode<()> {
45    /// Returns a new [`Builder`].
46    pub fn builder() -> Builder {
47        Builder::default()
48    }
49}
50
51impl Builder {
52    /// Sets the `app_state_bytes` to send the ABCI application upon chain initialization.
53    pub fn app_state(self, app_state: impl Into<Bytes>) -> Self {
54        let Self {
55            app_state: prev, ..
56        } = self;
57
58        // Log a warning if we are about to overwrite a previous value.
59        if let Some(prev) = prev {
60            tracing::warn!(
61                ?prev,
62                "builder overwriting a previously set `app_state`, this may be a bug!"
63            );
64        }
65
66        Self {
67            app_state: Some(app_state.into()),
68            ..self
69        }
70    }
71
72    /// Generates a single set of validator keys.
73    pub fn single_validator(self) -> Self {
74        let Self { keyring: prev, .. } = self;
75
76        // Log a warning if we are about to overwrite any existing keys.
77        if !prev.is_empty() {
78            tracing::warn!(
79                count = %prev.len(),
80                "builder overwriting entries in keyring, this may be a bug!"
81            );
82        }
83
84        // Generate a key and place it in the keyring.
85        let mut keyring = Keyring::new();
86        if self.keys.len() >= 1 {
87            Self::add_key(&mut keyring, self.keys[0].clone());
88        } else {
89            let key = Self::generate_key();
90            Self::add_key(&mut keyring, key);
91        }
92
93        Self { keyring, ..self }
94    }
95
96    /// Generates a pair of validator keys.
97    pub fn two_validators(self) -> Self {
98        let Self { keyring: prev, .. } = self;
99
100        // Log a warning if we are about to overwrite any existing keys.
101        if !prev.is_empty() {
102            tracing::warn!(
103                count = %prev.len(),
104                "builder overwriting entries in keyring, this may be a bug!"
105            );
106        }
107
108        // Generate two keys and place them in the keyring.
109        let mut keyring = Keyring::new();
110        if self.keys.len() >= 2 {
111            Self::add_key(&mut keyring, self.keys[0].clone());
112            Self::add_key(&mut keyring, self.keys[1].clone());
113        } else {
114            let key = Self::generate_key();
115            Self::add_key(&mut keyring, key);
116            let key = Self::generate_key();
117            Self::add_key(&mut keyring, key);
118        }
119
120        Self { keyring, ..self }
121    }
122
123    /// Generates consensus keys.
124    fn generate_key() -> (SigningKey, VerificationKey) {
125        let sk = ed25519_consensus::SigningKey::new(rand_core::OsRng);
126        let vk = sk.verification_key();
127        tracing::trace!(verification_key = ?vk, "generated consensus key");
128
129        (sk, vk)
130    }
131
132    /// Places keys in the provided keyring.
133    fn add_key(keyring: &mut Keyring, key: (SigningKey, VerificationKey)) {
134        let (sk, vk) = key;
135        keyring.insert(vk, sk);
136    }
137
138    /// Sets a callback that will be invoked when a new block is constructed.
139    pub fn on_block<F>(self, f: F) -> Self
140    where
141        F: FnMut(tendermint::Block) + Send + Sync + 'static,
142    {
143        Self {
144            on_block: Some(Box::new(f)),
145            ..self
146        }
147    }
148
149    /// Sets a callback that will be invoked when a block is committed, to increment
150    /// the timestamp.
151    pub fn ts_callback<F>(self, f: F) -> Self
152    where
153        F: Fn(Time) -> Time + Send + Sync + 'static,
154    {
155        Self {
156            ts_callback: Some(Box::new(f)),
157            ..self
158        }
159    }
160
161    /// Sets the starting time for the test node. If not called,
162    /// the current timestamp will be used.
163    pub fn with_initial_timestamp(self, initial_time: Time) -> Self {
164        Self {
165            initial_timestamp: Some(initial_time),
166            ..self
167        }
168    }
169
170    /// Sets the keys used by validators.
171    pub fn with_keys(self, keys: Vec<(SigningKey, VerificationKey)>) -> Self {
172        let Self {
173            keyring: ref prev, ..
174        } = self;
175
176        // Fail with a warning if we are about to overwrite any existing keys.
177        if !prev.is_empty() {
178            panic!("with_keys should be called prior to constructing the keyring");
179        }
180
181        Self { keys: keys, ..self }
182    }
183
184    /// Add the provided Tendermint [`Genesis`] to the builder.
185    ///
186    /// This will override other configurations and hardcode the genesis data.
187    pub fn with_tendermint_genesis(self, genesis: Genesis) -> Self {
188        let Self { .. } = &self;
189        Self {
190            app_state: Some(serde_json::to_vec(&genesis.app_state).unwrap().into()),
191            initial_timestamp: Some(genesis.genesis_time),
192            hardcoded_genesis: Some(genesis),
193            ..self
194        }
195    }
196}
penumbra_sdk_mock_consensus/builder.rs

penumbra_sdk_mock_consensus/
builder.rs
