test: add new test functions for engine tree (#9869)

crates/chain-state/src/test_utils.rs

@@ -2,39 +2,165 @@ use crate::{
     in_memory::ExecutedBlock, CanonStateNotification, CanonStateNotifications,
     CanonStateSubscriptions,
 };
-use rand::Rng;
+use rand::{thread_rng, Rng};
+use reth_chainspec::ChainSpec;
 use reth_execution_types::{Chain, ExecutionOutcome};
 use reth_primitives::{
-    Address, Block, BlockNumber, Receipts, Requests, SealedBlockWithSenders, TransactionSigned,
-    B256,
+    constants::{EIP1559_INITIAL_BASE_FEE, EMPTY_ROOT_HASH},
+    proofs::{calculate_receipt_root, calculate_transaction_root},
+    Address, BlockNumber, Header, Receipt, Receipts, Requests, SealedBlock, SealedBlockWithSenders,
+    Signature, Transaction, TransactionSigned, TransactionSignedEcRecovered, TxEip1559, B256, U256,
 };
-use reth_trie::{updates::TrieUpdates, HashedPostState};
-use revm::db::BundleState;
+use reth_trie::{root::state_root_unhashed, updates::TrieUpdates, HashedPostState};
+use revm::{db::BundleState, primitives::AccountInfo};
 use std::{
+    collections::HashMap,
     ops::Range,
     sync::{Arc, Mutex},
 };
 use tokio::sync::broadcast::{self, Sender};
 
+/// Generates a random `SealedBlockWithSenders`.
+pub fn generate_random_block(
+    number: BlockNumber,
+    parent_hash: B256,
+    chain_spec: &ChainSpec,
+    signer: Address,
+    signer_balance: &mut U256,
+    nonce: &mut u64,
+) -> SealedBlockWithSenders {
+    let mut rng = thread_rng();
+
+    let single_tx_cost = U256::from(EIP1559_INITIAL_BASE_FEE * 21_000);
+    let mock_tx = |nonce: u64| -> TransactionSignedEcRecovered {
+        TransactionSigned::from_transaction_and_signature(
+            Transaction::Eip1559(TxEip1559 {
+                chain_id: chain_spec.chain.id(),
+                nonce,
+                gas_limit: 21_000,
+                to: Address::random().into(),
+                max_fee_per_gas: EIP1559_INITIAL_BASE_FEE as u128,
+                max_priority_fee_per_gas: 1,
+                ..Default::default()
+            }),
+            Signature::default(),
+        )
+        .with_signer(signer)
+    };
+
+    let num_txs = rng.gen_range(0..5);
+    let signer_balance_decrease = single_tx_cost * U256::from(num_txs);
+    let transactions: Vec<TransactionSignedEcRecovered> = (0..num_txs)
+        .map(|_| {
+            let tx = mock_tx(*nonce);
+            *nonce += 1;
+            *signer_balance -= signer_balance_decrease;
+            tx
+        })
+        .collect();
+
+    let receipts = transactions
+        .iter()
+        .enumerate()
+        .map(|(idx, tx)| {
+            Receipt {
+                tx_type: tx.tx_type(),
+                success: true,
+                cumulative_gas_used: (idx as u64 + 1) * 21_000,
+                ..Default::default()
+            }
+            .with_bloom()
+        })
+        .collect::<Vec<_>>();
+
+    let initial_signer_balance = U256::from(10).pow(U256::from(18));
+
+    let header = Header {
+        number,
+        parent_hash,
+        gas_used: transactions.len() as u64 * 21_000,
+        gas_limit: chain_spec.max_gas_limit,
+        mix_hash: B256::random(),
+        base_fee_per_gas: Some(EIP1559_INITIAL_BASE_FEE),
+        transactions_root: calculate_transaction_root(&transactions),
+        receipts_root: calculate_receipt_root(&receipts),
+        beneficiary: Address::random(),
+        state_root: state_root_unhashed(HashMap::from([(
+            signer,
+            (
+                AccountInfo {
+                    balance: initial_signer_balance - signer_balance_decrease,
+                    nonce: num_txs,
+                    ..Default::default()
+                },
+                EMPTY_ROOT_HASH,
+            ),
+        )])),
+        ..Default::default()
+    };
+
+    let block = SealedBlock {
+        header: header.seal_slow(),
+        body: transactions.into_iter().map(|tx| tx.into_signed()).collect(),
+        ommers: Vec::new(),
+        withdrawals: None,
+        requests: None,
+    };
+
+    SealedBlockWithSenders::new(block, vec![signer; num_txs as usize]).unwrap()
+}
+
+/// Creates a fork chain with the given base block.
+pub fn create_fork(
+    base_block: &SealedBlock,
+    length: u64,
+    chain_spec: &ChainSpec,
+    signer: Address,
+    initial_signer_balance: U256,
+) -> Vec<SealedBlockWithSenders> {
+    let mut fork = Vec::with_capacity(length as usize);
+    let mut parent = base_block.clone();
+    let mut signer_balance = initial_signer_balance;
+    let mut nonce = 0;
+
+    for _ in 0..length {
+        let block = generate_random_block(
+            parent.number + 1,
+            parent.hash(),
+            chain_spec,
+            signer,
+            &mut signer_balance,
+            &mut nonce,
+        );
+        parent = block.block.clone();
+        fork.push(block);
+    }
+
+    fork
+}
+
 fn get_executed_block(
     block_number: BlockNumber,
     receipts: Receipts,
     parent_hash: B256,
 ) -> ExecutedBlock {
-    let mut block = Block::default();
-    let mut header = block.header.clone();
-    header.number = block_number;
-    header.parent_hash = parent_hash;
-    header.ommers_hash = B256::random();
-    block.header = header;
-    let tx = TransactionSigned::default();
-    block.body.push(tx);
-    let sealed = block.seal_slow();
-    let sender = Address::random();
-    let sealed_with_senders = SealedBlockWithSenders::new(sealed.clone(), vec![sender]).unwrap();
+    let chain_spec = ChainSpec::default();
+    let signer = Address::random();
+    let mut signer_balance = U256::from(1_000_000_000_000_000_000u64);
+    let mut nonce = 0;
+
+    let block_with_senders = generate_random_block(
+        block_number,
+        parent_hash,
+        &chain_spec,
+        signer,
+        &mut signer_balance,
+        &mut nonce,
+    );
+
     ExecutedBlock::new(
-        Arc::new(sealed),
-        Arc::new(sealed_with_senders.senders),
+        Arc::new(block_with_senders.block.clone()),
+        Arc::new(block_with_senders.senders),
         Arc::new(ExecutionOutcome::new(
             BundleState::default(),
             receipts,
@@ -64,7 +190,8 @@ pub fn get_executed_block_with_number(
 pub fn get_executed_blocks(range: Range<u64>) -> impl Iterator<Item = ExecutedBlock> {
     let mut parent_hash = B256::default();
     range.map(move |number| {
-        let block = get_executed_block_with_number(number, parent_hash);
+        let current_parent_hash = parent_hash;
+        let block = get_executed_block_with_number(number, current_parent_hash);
         parent_hash = block.block.hash();
         block
     })