reth_transaction_pool/test_utils/
pool.rs

Help
1//! Test helpers for mocking an entire pool.
2
3#![allow(dead_code)]
4
5use crate::{
6    pool::{txpool::TxPool, AddedTransaction},
7    test_utils::{MockOrdering, MockTransactionDistribution, MockTransactionFactory},
8    TransactionOrdering,
9};
10use alloy_primitives::{Address, U256};
11use rand::Rng;
12use std::{
13    collections::HashMap,
14    ops::{Deref, DerefMut},
15};
16
17/// A wrapped `TxPool` with additional helpers for testing
18pub(crate) struct MockPool<T: TransactionOrdering = MockOrdering> {
19    // The wrapped pool.
20    pool: TxPool<T>,
21}
22
23impl MockPool {
24    /// The total size of all subpools
25    fn total_subpool_size(&self) -> usize {
26        self.pool.pending().len() + self.pool.base_fee().len() + self.pool.queued().len()
27    }
28
29    /// Checks that all pool invariants hold.
30    fn enforce_invariants(&self) {
31        assert_eq!(
32            self.pool.len(),
33            self.total_subpool_size(),
34            "Tx in AllTransactions and sum(subpools) must match"
35        );
36    }
37}
38
39impl Default for MockPool {
40    fn default() -> Self {
41        Self { pool: TxPool::new(MockOrdering::default(), Default::default()) }
42    }
43}
44
45impl<T: TransactionOrdering> Deref for MockPool<T> {
46    type Target = TxPool<T>;
47
48    fn deref(&self) -> &Self::Target {
49        &self.pool
50    }
51}
52
53impl<T: TransactionOrdering> DerefMut for MockPool<T> {
54    fn deref_mut(&mut self) -> &mut Self::Target {
55        &mut self.pool
56    }
57}
58
59/// Simulates transaction execution.
60pub(crate) struct MockTransactionSimulator<R: Rng> {
61    /// The pending base fee
62    base_fee: u128,
63    /// Generator for transactions
64    tx_generator: MockTransactionDistribution,
65    /// represents the on chain balance of a sender.
66    balances: HashMap<Address, U256>,
67    /// represents the on chain nonce of a sender.
68    nonces: HashMap<Address, u64>,
69    /// A set of addresses to as senders.
70    senders: Vec<Address>,
71    /// What scenarios to execute.
72    scenarios: Vec<ScenarioType>,
73    /// All previous scenarios executed by a sender.
74    executed: HashMap<Address, ExecutedScenarios>,
75    /// "Validates" generated transactions.
76    validator: MockTransactionFactory,
77    /// The rng instance used to select senders and scenarios.
78    rng: R,
79}
80
81impl<R: Rng> MockTransactionSimulator<R> {
82    /// Returns a new mock instance
83    pub(crate) fn new(mut rng: R, config: MockSimulatorConfig) -> Self {
84        let senders = config.addresses(&mut rng);
85        Self {
86            base_fee: config.base_fee,
87            balances: senders.iter().copied().map(|a| (a, rng.random())).collect(),
88            nonces: senders.iter().copied().map(|a| (a, 0)).collect(),
89            senders,
90            scenarios: config.scenarios,
91            tx_generator: config.tx_generator,
92            executed: Default::default(),
93            validator: Default::default(),
94            rng,
95        }
96    }
97
98    /// Returns a random address from the senders set
99    fn rng_address(&mut self) -> Address {
100        let idx = self.rng.random_range(0..self.senders.len());
101        self.senders[idx]
102    }
103
104    /// Returns a random scenario from the scenario set
105    fn rng_scenario(&mut self) -> ScenarioType {
106        let idx = self.rng.random_range(0..self.scenarios.len());
107        self.scenarios[idx].clone()
108    }
109
110    /// Executes the next scenario and applies it to the pool
111    pub(crate) fn next(&mut self, pool: &mut MockPool) {
112        let sender = self.rng_address();
113        let scenario = self.rng_scenario();
114        let on_chain_nonce = self.nonces[&sender];
115        let on_chain_balance = self.balances[&sender];
116
117        match scenario {
118            ScenarioType::OnchainNonce => {
119                let tx = self
120                    .tx_generator
121                    .tx(on_chain_nonce, &mut self.rng)
122                    .with_gas_price(self.base_fee);
123                let valid_tx = self.validator.validated(tx);
124
125                let res = pool.add_transaction(valid_tx, on_chain_balance, on_chain_nonce).unwrap();
126
127                // TODO(mattsse): need a way expect based on the current state of the pool and tx
128                // settings
129
130                match res {
131                    AddedTransaction::Pending(_) => {}
132                    AddedTransaction::Parked { .. } => {
133                        panic!("expected pending")
134                    }
135                }
136
137                // TODO(mattsse): check subpools
138            }
139            ScenarioType::HigherNonce { .. } => {
140                unimplemented!()
141            }
142        }
143
144        // make sure everything is set
145        pool.enforce_invariants()
146    }
147}
148
149/// How to configure a new mock transaction stream
150pub(crate) struct MockSimulatorConfig {
151    /// How many senders to generate.
152    pub(crate) num_senders: usize,
153    /// Scenarios to test
154    pub(crate) scenarios: Vec<ScenarioType>,
155    /// The start base fee
156    pub(crate) base_fee: u128,
157    /// generator for transactions
158    pub(crate) tx_generator: MockTransactionDistribution,
159}
160
161impl MockSimulatorConfig {
162    /// Generates a set of random addresses
163    pub(crate) fn addresses(&self, rng: &mut impl rand::Rng) -> Vec<Address> {
164        std::iter::repeat_with(|| Address::random_with(rng)).take(self.num_senders).collect()
165    }
166}
167
168/// Represents
169#[derive(Debug, Clone)]
170#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
171pub(crate) enum ScenarioType {
172    OnchainNonce,
173    HigherNonce { skip: u64 },
174}
175
176/// The actual scenario, ready to be executed
177///
178/// A scenario produces one or more transactions and expects a certain Outcome.
179///
180/// An executed scenario can affect previous executed transactions
181#[derive(Debug, Clone)]
182#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
183pub(crate) enum Scenario {
184    /// Send a tx with the same nonce as on chain.
185    OnchainNonce { nonce: u64 },
186    /// Send a tx with a higher nonce that what the sender has on chain
187    HigherNonce { onchain: u64, nonce: u64 },
188    Multi {
189        // Execute multiple test scenarios
190        scenario: Vec<Scenario>,
191    },
192}
193
194/// Represents an executed scenario
195#[derive(Debug, Clone)]
196#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
197pub(crate) struct ExecutedScenario {
198    /// balance at the time of execution
199    balance: U256,
200    /// nonce at the time of execution
201    nonce: u64,
202    /// The executed scenario
203    scenario: Scenario,
204}
205
206/// All executed scenarios by a sender
207#[derive(Debug, Clone)]
208#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
209pub(crate) struct ExecutedScenarios {
210    sender: Address,
211    scenarios: Vec<ExecutedScenario>,
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217    use crate::test_utils::{MockFeeRange, MockTransactionRatio};
218
219    #[test]
220    fn test_on_chain_nonce_scenario() {
221        let transaction_ratio = MockTransactionRatio {
222            legacy_pct: 30,
223            dynamic_fee_pct: 70,
224            access_list_pct: 0,
225            blob_pct: 0,
226        };
227
228        let fee_ranges = MockFeeRange {
229            gas_price: (10u128..100).try_into().unwrap(),
230            priority_fee: (10u128..100).try_into().unwrap(),
231            max_fee: (100u128..110).try_into().unwrap(),
232            max_fee_blob: (1u128..100).try_into().unwrap(),
233        };
234
235        let config = MockSimulatorConfig {
236            num_senders: 10,
237            scenarios: vec![ScenarioType::OnchainNonce],
238            base_fee: 10,
239            tx_generator: MockTransactionDistribution::new(
240                transaction_ratio,
241                fee_ranges,
242                10..100,
243                10..100,
244            ),
245        };
246        let mut simulator = MockTransactionSimulator::new(rand::rng(), config);
247        let mut pool = MockPool::default();
248
249        simulator.next(&mut pool);
250    }
251}
reth_transaction_pool/test_utils/pool.rs

reth_transaction_pool/test_utils/
pool.rs
