reth_chain_state/
test_utils.rs

1use crate::{
2    in_memory::ExecutedBlockWithTrieUpdates, CanonStateNotification, CanonStateNotifications,
3    CanonStateSubscriptions, ExecutedTrieUpdates,
4};
5use alloy_consensus::{Header, SignableTransaction, TxEip1559, TxReceipt, EMPTY_ROOT_HASH};
6use alloy_eips::{
7    eip1559::{ETHEREUM_BLOCK_GAS_LIMIT_30M, INITIAL_BASE_FEE},
8    eip7685::Requests,
9};
10use alloy_primitives::{Address, BlockNumber, B256, U256};
11use alloy_signer::SignerSync;
12use alloy_signer_local::PrivateKeySigner;
13use core::marker::PhantomData;
14use rand::Rng;
15use reth_chainspec::{ChainSpec, EthereumHardfork, MIN_TRANSACTION_GAS};
16use reth_ethereum_primitives::{
17    Block, BlockBody, EthPrimitives, Receipt, Transaction, TransactionSigned,
18};
19use reth_execution_types::{Chain, ExecutionOutcome};
20use reth_primitives_traits::{
21    proofs::{calculate_receipt_root, calculate_transaction_root, calculate_withdrawals_root},
22    Account, NodePrimitives, Recovered, RecoveredBlock, SealedBlock, SealedHeader,
23    SignedTransaction,
24};
25use reth_storage_api::NodePrimitivesProvider;
26use reth_trie::{root::state_root_unhashed, HashedPostState};
27use revm_database::BundleState;
28use revm_state::AccountInfo;
29use std::{
30    collections::HashMap,
31    ops::Range,
32    sync::{Arc, Mutex},
33};
34use tokio::sync::broadcast::{self, Sender};
35
36/// Functionality to build blocks for tests and help with assertions about
37/// their execution.
38#[derive(Debug)]
39pub struct TestBlockBuilder<N: NodePrimitives = EthPrimitives> {
40    /// The account that signs all the block's transactions.
41    pub signer: Address,
42    /// Private key for signing.
43    pub signer_pk: PrivateKeySigner,
44    /// Keeps track of signer's account info after execution, will be updated in
45    /// methods related to block execution.
46    pub signer_execute_account_info: AccountInfo,
47    /// Keeps track of signer's nonce, will be updated in methods related
48    /// to block execution.
49    pub signer_build_account_info: AccountInfo,
50    /// Chain spec of the blocks generated by this builder
51    pub chain_spec: ChainSpec,
52    _prims: PhantomData<N>,
53}
54
55impl<N: NodePrimitives> Default for TestBlockBuilder<N> {
56    fn default() -> Self {
57        let initial_account_info = AccountInfo::from_balance(U256::from(10).pow(U256::from(18)));
58        let signer_pk = PrivateKeySigner::random();
59        let signer = signer_pk.address();
60        Self {
61            chain_spec: ChainSpec::default(),
62            signer,
63            signer_pk,
64            signer_execute_account_info: initial_account_info.clone(),
65            signer_build_account_info: initial_account_info,
66            _prims: PhantomData,
67        }
68    }
69}
70
71impl<N: NodePrimitives> TestBlockBuilder<N> {
72    /// Signer pk setter.
73    pub fn with_signer_pk(mut self, signer_pk: PrivateKeySigner) -> Self {
74        self.signer = signer_pk.address();
75        self.signer_pk = signer_pk;
76
77        self
78    }
79
80    /// Chainspec setter.
81    pub fn with_chain_spec(mut self, chain_spec: ChainSpec) -> Self {
82        self.chain_spec = chain_spec;
83        self
84    }
85
86    /// Gas cost of a single transaction generated by the block builder.
87    pub fn single_tx_cost() -> U256 {
88        U256::from(INITIAL_BASE_FEE * MIN_TRANSACTION_GAS)
89    }
90
91    /// Generates a random [`RecoveredBlock`].
92    pub fn generate_random_block(
93        &mut self,
94        number: BlockNumber,
95        parent_hash: B256,
96    ) -> RecoveredBlock<reth_ethereum_primitives::Block> {
97        let mut rng = rand::rng();
98
99        let mock_tx = |nonce: u64| -> Recovered<_> {
100            let tx = Transaction::Eip1559(TxEip1559 {
101                chain_id: self.chain_spec.chain.id(),
102                nonce,
103                gas_limit: MIN_TRANSACTION_GAS,
104                to: Address::random().into(),
105                max_fee_per_gas: INITIAL_BASE_FEE as u128,
106                max_priority_fee_per_gas: 1,
107                ..Default::default()
108            });
109            let signature_hash = tx.signature_hash();
110            let signature = self.signer_pk.sign_hash_sync(&signature_hash).unwrap();
111
112            TransactionSigned::new_unhashed(tx, signature).with_signer(self.signer)
113        };
114
115        let num_txs = rng.random_range(0..5);
116        let signer_balance_decrease = Self::single_tx_cost() * U256::from(num_txs);
117        let transactions: Vec<Recovered<_>> = (0..num_txs)
118            .map(|_| {
119                let tx = mock_tx(self.signer_build_account_info.nonce);
120                self.signer_build_account_info.nonce += 1;
121                self.signer_build_account_info.balance -= signer_balance_decrease;
122                tx
123            })
124            .collect();
125
126        let receipts = transactions
127            .iter()
128            .enumerate()
129            .map(|(idx, tx)| {
130                Receipt {
131                    tx_type: tx.tx_type(),
132                    success: true,
133                    cumulative_gas_used: (idx as u64 + 1) * MIN_TRANSACTION_GAS,
134                    ..Default::default()
135                }
136                .into_with_bloom()
137            })
138            .collect::<Vec<_>>();
139
140        let initial_signer_balance = U256::from(10).pow(U256::from(18));
141
142        let header = Header {
143            number,
144            parent_hash,
145            gas_used: transactions.len() as u64 * MIN_TRANSACTION_GAS,
146            mix_hash: B256::random(),
147            gas_limit: ETHEREUM_BLOCK_GAS_LIMIT_30M,
148            base_fee_per_gas: Some(INITIAL_BASE_FEE),
149            transactions_root: calculate_transaction_root(
150                &transactions.clone().into_iter().map(|tx| tx.into_inner()).collect::<Vec<_>>(),
151            ),
152            receipts_root: calculate_receipt_root(&receipts),
153            beneficiary: Address::random(),
154            state_root: state_root_unhashed(HashMap::from([(
155                self.signer,
156                Account {
157                    balance: initial_signer_balance - signer_balance_decrease,
158                    nonce: num_txs,
159                    ..Default::default()
160                }
161                .into_trie_account(EMPTY_ROOT_HASH),
162            )])),
163            // use the number as the timestamp so it is monotonically increasing
164            timestamp: number +
165                EthereumHardfork::Cancun.activation_timestamp(self.chain_spec.chain).unwrap(),
166            withdrawals_root: Some(calculate_withdrawals_root(&[])),
167            blob_gas_used: Some(0),
168            excess_blob_gas: Some(0),
169            parent_beacon_block_root: Some(B256::random()),
170            ..Default::default()
171        };
172
173        let block = SealedBlock::from_sealed_parts(
174            SealedHeader::seal_slow(header),
175            BlockBody {
176                transactions: transactions.into_iter().map(|tx| tx.into_inner()).collect(),
177                ommers: Vec::new(),
178                withdrawals: Some(vec![].into()),
179            },
180        );
181
182        RecoveredBlock::try_recover_sealed_with_senders(block, vec![self.signer; num_txs as usize])
183            .unwrap()
184    }
185
186    /// Creates a fork chain with the given base block.
187    pub fn create_fork(
188        &mut self,
189        base_block: &SealedBlock<Block>,
190        length: u64,
191    ) -> Vec<RecoveredBlock<Block>> {
192        let mut fork = Vec::with_capacity(length as usize);
193        let mut parent = base_block.clone();
194
195        for _ in 0..length {
196            let block = self.generate_random_block(parent.number + 1, parent.hash());
197            parent = block.clone_sealed_block();
198            fork.push(block);
199        }
200
201        fork
202    }
203
204    /// Gets an [`ExecutedBlockWithTrieUpdates`] with [`BlockNumber`], receipts and parent hash.
205    fn get_executed_block(
206        &mut self,
207        block_number: BlockNumber,
208        receipts: Vec<Vec<Receipt>>,
209        parent_hash: B256,
210    ) -> ExecutedBlockWithTrieUpdates {
211        let block_with_senders = self.generate_random_block(block_number, parent_hash);
212
213        let (block, senders) = block_with_senders.split_sealed();
214        ExecutedBlockWithTrieUpdates::new(
215            Arc::new(RecoveredBlock::new_sealed(block, senders)),
216            Arc::new(ExecutionOutcome::new(
217                BundleState::default(),
218                receipts,
219                block_number,
220                vec![Requests::default()],
221            )),
222            Arc::new(HashedPostState::default()),
223            ExecutedTrieUpdates::empty(),
224        )
225    }
226
227    /// Generates an [`ExecutedBlockWithTrieUpdates`] that includes the given receipts.
228    pub fn get_executed_block_with_receipts(
229        &mut self,
230        receipts: Vec<Vec<Receipt>>,
231        parent_hash: B256,
232    ) -> ExecutedBlockWithTrieUpdates {
233        let number = rand::rng().random::<u64>();
234        self.get_executed_block(number, receipts, parent_hash)
235    }
236
237    /// Generates an [`ExecutedBlockWithTrieUpdates`] with the given [`BlockNumber`].
238    pub fn get_executed_block_with_number(
239        &mut self,
240        block_number: BlockNumber,
241        parent_hash: B256,
242    ) -> ExecutedBlockWithTrieUpdates {
243        self.get_executed_block(block_number, vec![vec![]], parent_hash)
244    }
245
246    /// Generates a range of executed blocks with ascending block numbers.
247    pub fn get_executed_blocks(
248        &mut self,
249        range: Range<u64>,
250    ) -> impl Iterator<Item = ExecutedBlockWithTrieUpdates> + '_ {
251        let mut parent_hash = B256::default();
252        range.map(move |number| {
253            let current_parent_hash = parent_hash;
254            let block = self.get_executed_block_with_number(number, current_parent_hash);
255            parent_hash = block.recovered_block().hash();
256            block
257        })
258    }
259
260    /// Returns the execution outcome for a block created with this builder.
261    /// In order to properly include the bundle state, the signer balance is
262    /// updated.
263    pub fn get_execution_outcome(
264        &mut self,
265        block: RecoveredBlock<reth_ethereum_primitives::Block>,
266    ) -> ExecutionOutcome {
267        let num_txs = block.body().transactions.len() as u64;
268        let single_cost = Self::single_tx_cost();
269
270        let mut final_balance = self.signer_execute_account_info.balance;
271        for _ in 0..num_txs {
272            final_balance -= single_cost;
273        }
274
275        let final_nonce = self.signer_execute_account_info.nonce + num_txs;
276
277        let receipts = block
278            .body()
279            .transactions
280            .iter()
281            .enumerate()
282            .map(|(idx, tx)| Receipt {
283                tx_type: tx.tx_type(),
284                success: true,
285                cumulative_gas_used: (idx as u64 + 1) * MIN_TRANSACTION_GAS,
286                ..Default::default()
287            })
288            .collect::<Vec<_>>();
289
290        let bundle_state = BundleState::builder(block.number..=block.number)
291            .state_present_account_info(
292                self.signer,
293                AccountInfo { nonce: final_nonce, balance: final_balance, ..Default::default() },
294            )
295            .build();
296
297        self.signer_execute_account_info.balance = final_balance;
298        self.signer_execute_account_info.nonce = final_nonce;
299
300        let execution_outcome =
301            ExecutionOutcome::new(bundle_state, vec![vec![]], block.number, Vec::new());
302
303        execution_outcome.with_receipts(vec![receipts])
304    }
305}
306
307impl TestBlockBuilder {
308    /// Creates a `TestBlockBuilder` configured for Ethereum primitives.
309    pub fn eth() -> Self {
310        Self::default()
311    }
312}
313/// A test `ChainEventSubscriptions`
314#[derive(Clone, Debug, Default)]
315pub struct TestCanonStateSubscriptions<N: NodePrimitives = reth_ethereum_primitives::EthPrimitives>
316{
317    canon_notif_tx: Arc<Mutex<Vec<Sender<CanonStateNotification<N>>>>>,
318}
319
320impl TestCanonStateSubscriptions {
321    /// Adds new block commit to the queue that can be consumed with
322    /// [`TestCanonStateSubscriptions::subscribe_to_canonical_state`]
323    pub fn add_next_commit(&self, new: Arc<Chain>) {
324        let event = CanonStateNotification::Commit { new };
325        self.canon_notif_tx.lock().as_mut().unwrap().retain(|tx| tx.send(event.clone()).is_ok())
326    }
327
328    /// Adds reorg to the queue that can be consumed with
329    /// [`TestCanonStateSubscriptions::subscribe_to_canonical_state`]
330    pub fn add_next_reorg(&self, old: Arc<Chain>, new: Arc<Chain>) {
331        let event = CanonStateNotification::Reorg { old, new };
332        self.canon_notif_tx.lock().as_mut().unwrap().retain(|tx| tx.send(event.clone()).is_ok())
333    }
334}
335
336impl NodePrimitivesProvider for TestCanonStateSubscriptions {
337    type Primitives = EthPrimitives;
338}
339
340impl CanonStateSubscriptions for TestCanonStateSubscriptions {
341    /// Sets up a broadcast channel with a buffer size of 100.
342    fn subscribe_to_canonical_state(&self) -> CanonStateNotifications {
343        let (canon_notif_tx, canon_notif_rx) = broadcast::channel(100);
344        self.canon_notif_tx.lock().as_mut().unwrap().push(canon_notif_tx);
345
346        canon_notif_rx
347    }
348}
reth_chain_state/test_utils.rs

reth_chain_state/
test_utils.rs
