Skip to main content

reth_provider/test_utils/
mod.rs

1use crate::{
2    providers::{
3        NodeTypesForProvider, ProviderNodeTypes, RocksDBBuilder, StaticFileProvider,
4        StaticFileProviderBuilder,
5    },
6    HashingWriter, ProviderFactory, TrieWriter,
7};
8use alloy_primitives::B256;
9use reth_chainspec::{ChainSpec, ChainSpecBuilder, MAINNET};
10use reth_db::{mdbx::DatabaseArguments, test_utils::TempDatabase, DatabaseEnv};
11use reth_errors::ProviderResult;
12use reth_ethereum_engine_primitives::EthEngineTypes;
13use reth_node_types::NodeTypesWithDBAdapter;
14use reth_primitives_traits::{Account, StorageEntry};
15use reth_storage_api::StorageSettingsCache;
16use reth_trie::StateRoot;
17use reth_trie_db::DatabaseStateRoot;
18use std::sync::Arc;
19
20type DbStateRoot<'a, TX, A> = StateRoot<
21    reth_trie_db::DatabaseTrieCursorFactory<&'a TX, A>,
22    reth_trie_db::DatabaseHashedCursorFactory<&'a TX>,
23>;
24
25pub mod blocks;
26mod mock;
27mod noop;
28
29pub use mock::{ExtendedAccount, MockEthProvider};
30pub use noop::NoopProvider;
31pub use reth_chain_state::test_utils::TestCanonStateSubscriptions;
32
33/// Mock [`reth_node_types::NodeTypes`] for testing.
34pub type MockNodeTypes = reth_node_types::AnyNodeTypesWithEngine<
35    reth_ethereum_primitives::EthPrimitives,
36    reth_ethereum_engine_primitives::EthEngineTypes,
37    reth_chainspec::ChainSpec,
38    crate::EthStorage,
39    EthEngineTypes,
40>;
41
42/// Mock [`reth_node_types::NodeTypesWithDB`] for testing.
43pub type MockNodeTypesWithDB<DB = Arc<TempDatabase<DatabaseEnv>>> =
44    NodeTypesWithDBAdapter<MockNodeTypes, DB>;
45
46/// Creates test provider factory with mainnet chain spec.
47pub fn create_test_provider_factory() -> ProviderFactory<MockNodeTypesWithDB> {
48    create_test_provider_factory_with_chain_spec(MAINNET.clone())
49}
50
51/// Creates test provider factory with provided chain spec.
52pub fn create_test_provider_factory_with_chain_spec(
53    chain_spec: Arc<ChainSpec>,
54) -> ProviderFactory<MockNodeTypesWithDB> {
55    let genesis_block_number = chain_spec.genesis.number.unwrap_or_default();
56    create_test_provider_factory_with_node_types_and_genesis::<MockNodeTypes>(
57        chain_spec,
58        genesis_block_number,
59    )
60}
61
62/// Creates a test provider factory whose chain starts at `genesis_block_number`.
63pub fn create_test_provider_factory_with_genesis_block_number(
64    genesis_block_number: u64,
65) -> ProviderFactory<MockNodeTypesWithDB> {
66    let mut genesis = MAINNET.genesis.clone();
67    genesis.number = Some(genesis_block_number);
68    let chain_spec = Arc::new(ChainSpecBuilder::mainnet().genesis(genesis).build());
69    create_test_provider_factory_with_chain_spec(chain_spec)
70}
71
72/// Creates test provider factory with provided chain spec.
73pub fn create_test_provider_factory_with_node_types<N: NodeTypesForProvider>(
74    chain_spec: Arc<N::ChainSpec>,
75) -> ProviderFactory<NodeTypesWithDBAdapter<N, Arc<TempDatabase<DatabaseEnv>>>> {
76    create_test_provider_factory_with_node_types_and_genesis(chain_spec, 0)
77}
78
79fn create_test_provider_factory_with_node_types_and_genesis<N: NodeTypesForProvider>(
80    chain_spec: Arc<N::ChainSpec>,
81    genesis_block_number: u64,
82) -> ProviderFactory<NodeTypesWithDBAdapter<N, Arc<TempDatabase<DatabaseEnv>>>> {
83    // Create a single temp directory that contains all data dirs (db, static_files, rocksdb).
84    // TempDatabase will clean up the entire directory on drop.
85    let datadir_path = reth_db::test_utils::tempdir_path();
86
87    let static_files_path = datadir_path.join("static_files");
88    let rocksdb_path = datadir_path.join("rocksdb");
89
90    // Create static_files directory
91    std::fs::create_dir_all(&static_files_path).expect("failed to create static_files dir");
92
93    // Create database with the datadir path so TempDatabase cleans up everything on drop
94    let db = reth_db::test_utils::create_test_rw_db_with_datadir(&datadir_path);
95
96    ProviderFactory::new(
97        db,
98        chain_spec,
99        StaticFileProviderBuilder::read_write(static_files_path)
100            .with_genesis_block_number(genesis_block_number)
101            .build()
102            .expect("static file provider"),
103        RocksDBBuilder::new(&rocksdb_path)
104            .with_default_tables()
105            .build()
106            .expect("failed to create test RocksDB provider"),
107        reth_tasks::Runtime::test(),
108    )
109    .expect("failed to create test provider factory")
110}
111
112/// Creates test provider factory with provided chain spec and custom database arguments.
113///
114/// Same as [`create_test_provider_factory_with_chain_spec`] but allows overriding the default
115/// test database arguments (e.g. to increase the MDBX geometry for heavy benchmarks).
116pub fn create_test_provider_factory_with_chain_spec_and_db_args(
117    chain_spec: Arc<ChainSpec>,
118    db_args: DatabaseArguments,
119) -> ProviderFactory<MockNodeTypesWithDB> {
120    let datadir_path = reth_db::test_utils::tempdir_path();
121
122    let db_path = datadir_path.join("db");
123    let static_files_path = datadir_path.join("static_files");
124    let rocksdb_path = datadir_path.join("rocksdb");
125
126    std::fs::create_dir_all(&static_files_path).expect("failed to create static_files dir");
127
128    let db = reth_db::init_db(&db_path, db_args).expect("failed to init db");
129    let db = Arc::new(TempDatabase::new(db, datadir_path));
130
131    ProviderFactory::new(
132        db,
133        chain_spec,
134        StaticFileProvider::read_write(static_files_path).expect("static file provider"),
135        RocksDBBuilder::new(&rocksdb_path)
136            .with_default_tables()
137            .build()
138            .expect("failed to create test RocksDB provider"),
139        reth_tasks::Runtime::test(),
140    )
141    .expect("failed to create test provider factory")
142}
143
144/// Inserts the genesis alloc from the provided chain spec into the trie.
145pub fn insert_genesis<N: ProviderNodeTypes<ChainSpec = ChainSpec>>(
146    provider_factory: &ProviderFactory<N>,
147    chain_spec: Arc<N::ChainSpec>,
148) -> ProviderResult<B256> {
149    let provider = provider_factory.provider_rw()?;
150
151    // Hash accounts and insert them into hashing table.
152    let genesis = chain_spec.genesis();
153    let alloc_accounts =
154        genesis.alloc.iter().map(|(addr, account)| (*addr, Some(Account::from(account))));
155    provider.insert_account_for_hashing(alloc_accounts).unwrap();
156
157    let alloc_storage = genesis.alloc.clone().into_iter().filter_map(|(addr, account)| {
158        // Only return `Some` if there is storage.
159        account.storage.map(|storage| {
160            (
161                addr,
162                storage.into_iter().map(|(key, value)| StorageEntry { key, value: value.into() }),
163            )
164        })
165    });
166    provider.insert_storage_for_hashing(alloc_storage)?;
167
168    let (root, updates) = reth_trie_db::with_adapter!(provider, |A| {
169        DbStateRoot::<_, A>::from_tx(provider.tx_ref()).root_with_updates()?
170    });
171    provider.write_trie_updates(updates).unwrap();
172
173    provider.commit()?;
174
175    Ok(root)
176}