reth_chain_state/
test_utils.rs

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

reth_chain_state/
test_utils.rs
