clean up database file after testing

crates/consensus/beacon/src/engine/sync.rs

@@ -396,7 +396,7 @@ mod tests {
     use futures::poll;
     use reth_db::{
         mdbx::{Env, WriteMap},
-        test_utils::create_test_rw_db,
+        test_utils::{create_test_rw_db, TestTempDatabase},
     };
     use reth_interfaces::{p2p::either::EitherDownloader, test_utils::TestFullBlockClient};
     use reth_primitives::{
@@ -449,7 +449,10 @@ mod tests {
         }
 
         /// Builds the pipeline.
-        fn build(self, chain_spec: Arc<ChainSpec>) -> Pipeline<Arc<Env<WriteMap>>> {
+        fn build(
+            self,
+            chain_spec: Arc<ChainSpec>,
+        ) -> Pipeline<Arc<TestTempDatabase<Env<WriteMap>>>> {
             reth_tracing::init_test_tracing();
             let db = create_test_rw_db();
 

crates/consensus/beacon/src/engine/test_utils.rs

@@ -6,7 +6,11 @@ use crate::{
 use reth_blockchain_tree::{
     config::BlockchainTreeConfig, externals::TreeExternals, BlockchainTree, ShareableBlockchainTree,
 };
-use reth_db::{test_utils::create_test_rw_db, DatabaseEnv};
+use reth_db::{
+    test_utils::{create_test_rw_db, TestTempDatabase},
+    DatabaseEnv as DE,
+};
+type DatabaseEnv = TestTempDatabase<DE>;
 use reth_downloaders::{
     bodies::bodies::BodiesDownloaderBuilder,
     headers::reverse_headers::ReverseHeadersDownloaderBuilder,

crates/stages/Cargo.toml

@@ -52,7 +52,7 @@ num-traits = "0.2.15"
 
 [dev-dependencies]
 # reth
-reth-primitives = { workspace = true, features = ["arbitrary"] }
+reth-primitives = { workspace = true, features = ["test-utils", "arbitrary"] }
 reth-db = { workspace = true, features = ["test-utils", "mdbx"] }
 reth-interfaces = { workspace = true, features = ["test-utils"] }
 reth-downloaders = { path = "../net/downloaders" }

crates/stages/src/stages/bodies.rs

@@ -458,6 +458,7 @@ mod tests {
             database::Database,
             models::{StoredBlockBodyIndices, StoredBlockOmmers},
             tables,
+            test_utils::TestTempDatabase,
             transaction::{DbTx, DbTxMut},
             DatabaseEnv,
         };
@@ -740,15 +741,15 @@ mod tests {
         /// A [BodyDownloader] that is backed by an internal [HashMap] for testing.
         #[derive(Debug)]
         pub(crate) struct TestBodyDownloader {
-            db: Arc<DatabaseEnv>,
+            db: Arc<TestTempDatabase<DatabaseEnv>>,
             responses: HashMap<B256, BlockBody>,
             headers: VecDeque<SealedHeader>,
             batch_size: u64,
         }
 
         impl TestBodyDownloader {
             pub(crate) fn new(
-                db: Arc<DatabaseEnv>,
+                db: Arc<TestTempDatabase<DatabaseEnv>>,
                 responses: HashMap<B256, BlockBody>,
                 batch_size: u64,
             ) -> Self {

crates/stages/src/stages/mod.rs

@@ -69,7 +69,7 @@ mod tests {
     #[ignore]
     async fn test_prune() {
         let test_tx = TestTransaction::default();
-        let factory = Arc::new(ProviderFactory::new(test_tx.tx.as_ref(), MAINNET.clone()));
+        let factory = Arc::new(ProviderFactory::new(test_tx.tx.db(), MAINNET.clone()));
 
         let provider = factory.provider_rw().unwrap();
         let tip = 66;

crates/stages/src/test_utils/runner.rs

@@ -48,7 +48,7 @@ pub(crate) trait ExecuteStageTestRunner: StageTestRunner {
         let (tx, rx) = oneshot::channel();
         let (db, mut stage) = (self.tx().inner_raw(), self.stage());
         tokio::spawn(async move {
-            let factory = ProviderFactory::new(db.as_ref(), MAINNET.clone());
+            let factory = ProviderFactory::new(db.db(), MAINNET.clone());
             let provider = factory.provider_rw().unwrap();
 
             let result = stage.execute(&provider, input).await;
@@ -74,7 +74,7 @@ pub(crate) trait UnwindStageTestRunner: StageTestRunner {
         let (tx, rx) = oneshot::channel();
         let (db, mut stage) = (self.tx().inner_raw(), self.stage());
         tokio::spawn(async move {
-            let factory = ProviderFactory::new(db.as_ref(), MAINNET.clone());
+            let factory = ProviderFactory::new(db.db(), MAINNET.clone());
             let provider = factory.provider_rw().unwrap();
 
             let result = stage.unwind(&provider, input).await;

crates/stages/src/test_utils/test_db.rs

@@ -5,7 +5,7 @@ use reth_db::{
     models::{AccountBeforeTx, StoredBlockBodyIndices},
     table::{Table, TableRow},
     tables,
-    test_utils::{create_test_rw_db, create_test_rw_db_with_path},
+    test_utils::{create_test_rw_db, create_test_rw_db_with_path, TestTempDatabase},
     transaction::{DbTx, DbTxGAT, DbTxMut, DbTxMutGAT},
     DatabaseEnv, DatabaseError as DbError,
 };
@@ -33,9 +33,9 @@ use std::{
 #[derive(Debug)]
 pub struct TestTransaction {
     /// DB
-    pub tx: Arc<DatabaseEnv>,
+    pub tx: Arc<TestTempDatabase<DatabaseEnv>>,
     pub path: Option<PathBuf>,
-    pub factory: ProviderFactory<Arc<DatabaseEnv>>,
+    pub factory: ProviderFactory<Arc<TestTempDatabase<DatabaseEnv>>>,
 }
 
 impl Default for TestTransaction {
@@ -57,17 +57,17 @@ impl TestTransaction {
     }
 
     /// Return a database wrapped in [DatabaseProviderRW].
-    pub fn inner_rw(&self) -> DatabaseProviderRW<'_, Arc<DatabaseEnv>> {
+    pub fn inner_rw(&self) -> DatabaseProviderRW<'_, Arc<TestTempDatabase<DatabaseEnv>>> {
         self.factory.provider_rw().expect("failed to create db container")
     }
 
     /// Return a database wrapped in [DatabaseProviderRO].
-    pub fn inner(&self) -> DatabaseProviderRO<'_, Arc<DatabaseEnv>> {
+    pub fn inner(&self) -> DatabaseProviderRO<'_, Arc<TestTempDatabase<DatabaseEnv>>> {
         self.factory.provider().expect("failed to create db container")
     }
 
     /// Get a pointer to an internal database.
-    pub fn inner_raw(&self) -> Arc<DatabaseEnv> {
+    pub fn inner_raw(&self) -> Arc<TestTempDatabase<DatabaseEnv>> {
         self.tx.clone()
     }
 

crates/storage/db/benches/hash_keys.rs

@@ -86,6 +86,7 @@ where
             // Reset DB
             let _ = fs::remove_dir_all(bench_db_path);
             let db = Arc::try_unwrap(create_test_rw_db_with_path(bench_db_path)).unwrap();
+            let db = db.into_inner_db();
 
             let mut unsorted_input = unsorted_input.clone();
             if scenario_str == "append_all" {

crates/storage/db/benches/utils.rs

@@ -72,5 +72,5 @@ where
         tx.inner.commit().unwrap();
     }
 
-    db
+    db.into_inner_db()
 }

crates/storage/db/src/lib.rs

@@ -160,7 +160,8 @@ pub fn open_db(path: &Path, log_level: Option<LogLevel>) -> eyre::Result<Databas
 #[cfg(any(test, feature = "test-utils"))]
 pub mod test_utils {
     use super::*;
-    use std::sync::Arc;
+    use crate::database::{Database, DatabaseGAT};
+    use std::{path::PathBuf, sync::Arc};
 
     /// Error during database open
     pub const ERROR_DB_OPEN: &str = "Not able to open the database file.";
@@ -171,26 +172,76 @@ pub mod test_utils {
     /// Error during tempdir creation
     pub const ERROR_TEMPDIR: &str = "Not able to create a temporary directory.";
 
+    /// A database will delete the db dir when dropped.
+    #[derive(Debug)]
+    pub struct TestTempDatabase<DB> {
+        db: Option<DB>,
+        path: PathBuf,
+    }
+
+    impl<DB> Drop for TestTempDatabase<DB> {
+        fn drop(&mut self) {
+            if let Some(db) = self.db.take() {
+                drop(db);
+                let _ = std::fs::remove_dir_all(&self.path);
+            }
+        }
+    }
+
+    impl<DB> TestTempDatabase<DB> {
+        /// retunrs the ref of inner db
+        pub fn db(&self) -> &DB {
+            self.db.as_ref().unwrap()
+        }
+
+        /// retunrs the inner db
+        pub fn into_inner_db(mut self) -> DB {
+            self.db.take().unwrap() // take out db to avoid clean path in drop fn
+        }
+    }
+
+    impl<'a, DB: Database> DatabaseGAT<'a> for TestTempDatabase<DB> {
+        type TX = <DB as DatabaseGAT<'a>>::TX;
+        type TXMut = <DB as DatabaseGAT<'a>>::TXMut;
+    }
+
+    impl<DB: Database> Database for TestTempDatabase<DB> {
+        fn tx(&self) -> Result<<Self as DatabaseGAT<'_>>::TX, DatabaseError> {
+            self.db().tx()
+        }
+
+        fn tx_mut(&self) -> Result<<Self as DatabaseGAT<'_>>::TXMut, DatabaseError> {
+            self.db().tx_mut()
+        }
+    }
+
     /// Create read/write database for testing
-    pub fn create_test_rw_db() -> Arc<DatabaseEnv> {
-        Arc::new(
-            init_db(tempfile::TempDir::new().expect(ERROR_TEMPDIR).into_path(), None)
-                .expect(ERROR_DB_CREATION),
-        )
+    pub fn create_test_rw_db() -> Arc<TestTempDatabase<DatabaseEnv>> {
+        let path = tempfile::TempDir::new().expect(ERROR_TEMPDIR).into_path();
+        let emsg = format!("{}: {:?}", ERROR_DB_CREATION, path);
+
+        let db = init_db(&path, None).expect(&emsg);
+
+        Arc::new(TestTempDatabase { db: Some(db), path })
     }
 
     /// Create read/write database for testing
-    pub fn create_test_rw_db_with_path<P: AsRef<Path>>(path: P) -> Arc<DatabaseEnv> {
-        Arc::new(init_db(path.as_ref(), None).expect(ERROR_DB_CREATION))
+    pub fn create_test_rw_db_with_path<P: AsRef<Path>>(
+        path: P,
+    ) -> Arc<TestTempDatabase<DatabaseEnv>> {
+        let path = path.as_ref().to_path_buf();
+        let db = init_db(path.as_path(), None).expect(ERROR_DB_CREATION);
+        Arc::new(TestTempDatabase { db: Some(db), path })
     }
 
     /// Create read only database for testing
-    pub fn create_test_ro_db() -> Arc<DatabaseEnvRO> {
+    pub fn create_test_ro_db() -> Arc<TestTempDatabase<DatabaseEnvRO>> {
         let path = tempfile::TempDir::new().expect(ERROR_TEMPDIR).into_path();
         {
             init_db(path.as_path(), None).expect(ERROR_DB_CREATION);
         }
-        Arc::new(open_db_read_only(path.as_path(), None).expect(ERROR_DB_OPEN))
+        let db = open_db_read_only(path.as_path(), None).expect(ERROR_DB_OPEN);
+        Arc::new(TestTempDatabase { db: Some(db), path })
     }
 }
 

crates/trie/src/trie.rs

@@ -1138,7 +1138,7 @@ mod tests {
     #[test]
     fn account_trie_around_extension_node() {
         let db = create_test_rw_db();
-        let factory = ProviderFactory::new(db.as_ref(), MAINNET.clone());
+        let factory = ProviderFactory::new(db.db(), MAINNET.clone());
         let tx = factory.provider_rw().unwrap();
 
         let expected = extension_node_trie(&tx);
@@ -1164,7 +1164,7 @@ mod tests {
 
     fn account_trie_around_extension_node_with_dbtrie() {
         let db = create_test_rw_db();
-        let factory = ProviderFactory::new(db.as_ref(), MAINNET.clone());
+        let factory = ProviderFactory::new(db.db(), MAINNET.clone());
         let tx = factory.provider_rw().unwrap();
 
         let expected = extension_node_trie(&tx);
@@ -1227,7 +1227,7 @@ mod tests {
     #[test]
     fn storage_trie_around_extension_node() {
         let db = create_test_rw_db();
-        let factory = ProviderFactory::new(db.as_ref(), MAINNET.clone());
+        let factory = ProviderFactory::new(db.db(), MAINNET.clone());
         let tx = factory.provider_rw().unwrap();
 
         let hashed_address = B256::random();