reth_chain_state/
test_utils.rs

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