Skip to main content

reth_provider/test_utils/
mod.rs

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