reth_chain_state/
test_utils.rs

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

reth_chain_state/
test_utils.rs
