Refactor tests (#332)

firewood/tests/common/mod.rs

@@ -0,0 +1,72 @@
+// Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE.md for licensing terms.
+
+use std::{env::temp_dir, fs::remove_dir_all, ops::Deref, path::PathBuf};
+
+use firewood::db::{Db, DbConfig};
+use typed_builder::TypedBuilder;
+
+#[derive(Clone, Debug, TypedBuilder)]
+pub struct TestDbCreator {
+    #[builder(setter(into))]
+    test_name: String,
+    #[builder(default, setter(into))]
+    path: Option<PathBuf>,
+    #[builder(default = DbConfig::builder().truncate(true).build())]
+    cfg: DbConfig,
+}
+
+pub struct TestDb {
+    creator: TestDbCreator,
+    preserve_on_drop: bool,
+    db: Db,
+}
+
+impl TestDbCreator {
+    pub async fn create(self) -> TestDb {
+        let path = self.path.clone().unwrap_or_else(|| {
+            let mut path: PathBuf = std::env::var_os("CARGO_TARGET_DIR")
+                .unwrap_or(temp_dir().into())
+                .into();
+            if path.join("tmp").is_dir() {
+                path.push("tmp");
+            }
+            path.join(&self.test_name)
+        });
+        let mut creator = self.clone();
+        creator.path = path.clone().into();
+        let db = Db::new(&path, &self.cfg).await.unwrap();
+        TestDb {
+            creator,
+            db,
+            preserve_on_drop: false,
+        }
+    }
+}
+
+impl Deref for TestDb {
+    type Target = Db;
+
+    fn deref(&self) -> &Self::Target {
+        &self.db
+    }
+}
+
+impl TestDb {
+    /// reopen the database, consuming the old TestDb and giving you a new one
+    pub async fn reopen(mut self) -> Self {
+        let mut creator = self.creator.clone();
+        self.preserve_on_drop = true;
+        drop(self);
+        creator.cfg.truncate = false;
+        creator.create().await
+    }
+}
+
+impl Drop for TestDb {
+    fn drop(&mut self) {
+        if !self.preserve_on_drop {
+            remove_dir_all(self.creator.path.as_ref().unwrap()).unwrap();
+        }
+    }
+}

firewood/tests/db.rs

@@ -2,13 +2,16 @@
 // See the file LICENSE.md for licensing terms.
 
 use firewood::{
-    db::{Db, DbConfig, WalConfig},
+    db::{DbConfig, WalConfig},
     v2::api::{self, BatchOp, Db as _, DbView, Proposal},
 };
 use tokio::task::block_in_place;
 
 use std::{collections::VecDeque, env::temp_dir, path::PathBuf, sync::Arc};
 
+mod common;
+use common::TestDbCreator;
+
 // TODO: use a trait
 macro_rules! kv_dump {
     ($e: ident) => {{
@@ -35,14 +38,13 @@ async fn test_basic_metrics() {
                 .build(),
         );
 
-    let mut tmpdir: PathBuf = std::env::var_os("CARGO_TARGET_DIR")
-        .unwrap_or(temp_dir().into())
-        .into();
-    tmpdir.push("/tmp/test_basic_metrics");
+    let db = TestDbCreator::builder()
+        .cfg(cfg.build())
+        .test_name("test_basic_metrics")
+        .build()
+        .create()
+        .await;
 
-    let db = firewood::db::Db::new(tmpdir, &cfg.truncate(true).build())
-        .await
-        .unwrap();
     // let metrics = db.metrics();
     // TODO: kv_get is no longer a valid metric, and DbRev has no access to Db.metrics (yet)
     //assert_eq!(metrics.kv_get.hit_count.get(), 0);
@@ -82,7 +84,8 @@ async fn test_revisions() {
                 .block_nbit(8)
                 .max_revisions(10)
                 .build(),
-        );
+        )
+        .build();
 
     let rng = std::cell::RefCell::new(StdRng::seed_from_u64(42));
     let max_len0 = 8;
@@ -102,15 +105,13 @@ async fn test_revisions() {
         key
     };
 
-    let mut tmpdir: PathBuf = std::env::var_os("CARGO_TARGET_DIR")
-        .unwrap_or(temp_dir().into())
-        .into();
-    tmpdir.push("/tmp/test_revisions");
-
     for i in 0..10 {
-        let db = Db::new(&tmpdir, &cfg.clone().truncate(true).build())
-            .await
-            .unwrap();
+        let db = TestDbCreator::builder()
+            .cfg(cfg.clone())
+            .test_name("test_revisions")
+            .build()
+            .create()
+            .await;
         let mut dumped = VecDeque::new();
         let mut hashes: VecDeque<api::HashKey> = VecDeque::new();
         for _ in 0..10 {
@@ -142,10 +143,8 @@ async fn test_revisions() {
                 assert_eq!(kv_dump!(rev), *dump, "not the same: Pass {i}");
             }
         }
-        drop(db);
-        let db = Db::new(tmpdir.clone(), &cfg.clone().truncate(false).build())
-            .await
-            .unwrap();
+
+        let db = db.reopen().await;
         for (dump, hash) in dumped.iter().zip(hashes.iter().cloned()) {
             let rev = db.revision(hash).await.unwrap();
             rev.root_hash().await.unwrap();
@@ -247,16 +246,16 @@ macro_rules! assert_val {
 
 #[tokio::test(flavor = "multi_thread")]
 async fn db_proposal() -> Result<(), api::Error> {
-    let cfg = DbConfig::builder().wal(WalConfig::builder().max_revisions(10).build());
-
-    let mut tmpdir: PathBuf = std::env::var_os("CARGO_TARGET_DIR")
-        .unwrap_or(temp_dir().into())
-        .into();
-    tmpdir.push("/tmp/test_db_proposal");
-
-    let db = firewood::db::Db::new(tmpdir, &cfg.clone().truncate(true).build())
-        .await
-        .expect("db initiation should succeed");
+    let cfg = DbConfig::builder()
+        .wal(WalConfig::builder().max_revisions(10).build())
+        .build();
+
+    let db = TestDbCreator::builder()
+        .cfg(cfg)
+        .test_name("db_proposal")
+        .build()
+        .create()
+        .await;
 
     let batch = vec![
         BatchOp::Put {

firewood/tests/v2api.rs

@@ -1,17 +1,26 @@
 // Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE.md for licensing terms.
 
-use std::{path::PathBuf, sync::Arc};
+use std::sync::Arc;
 
 use firewood::{
-    db::{BatchOp, Db as PersistedDb, DbConfig},
+    db::{BatchOp, DbConfig},
     v2::api::{Db, DbView, Proposal},
 };
 
+pub mod common;
+use common::TestDbCreator;
+
 #[tokio::test(flavor = "multi_thread")]
 async fn smoke() -> Result<(), Box<dyn std::error::Error>> {
     let cfg = DbConfig::builder().truncate(true).build();
-    let db = Arc::new(testdb(cfg).await?);
+    let db = TestDbCreator::builder()
+        .cfg(cfg)
+        .test_name("smoke")
+        .build()
+        .create()
+        .await;
+
     let empty_hash = db.root_hash().await?;
     assert_ne!(empty_hash, [0; 32]);
 
@@ -40,14 +49,3 @@ async fn smoke() -> Result<(), Box<dyn std::error::Error>> {
 
     Ok(())
 }
-
-async fn testdb(cfg: DbConfig) -> Result<impl firewood::v2::api::Db, firewood::v2::api::Error> {
-    let tmpdbpath = tmp_dir().join("testdb");
-    PersistedDb::new(tmpdbpath, &cfg).await
-}
-
-fn tmp_dir() -> PathBuf {
-    option_env!("CARGO_TARGET_TMPDIR")
-        .map(PathBuf::from)
-        .unwrap_or(std::env::temp_dir())
-}