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