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

reth_transaction_pool/test_utils/
pool.rs
