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

reth_chain_state/
test_utils.rs
