From bdefd9af30bb5b3989a94532e7e034f152c1b4a3 Mon Sep 17 00:00:00 2001
From: Chen Chen <mail@roychan.org>
Date: Thu, 1 Aug 2024 13:48:35 -0500
Subject: [PATCH] stores: add rocksdb tests

---
 Cargo.toml            |  3 +++
 src/bench.rs          | 14 ++++++++++++++
 src/server.rs         | 32 ++++++++++++++++++++++++++++++++
 src/stores.rs         | 11 +++++++++++
 src/stores/rocksdb.rs |  2 +-
 5 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index a6fdb40..e198ff7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -42,6 +42,9 @@ serde = { version = "1.0.204", features = ["derive"] }
 toml = "0.8.19"
 zipf = "7.0.1"
 
+[dev-dependencies]
+tempfile = "3.10.1"
+
 [features]
 chashmap = ["dep:chashmap"]
 contrie = ["dep:contrie"]
diff --git a/src/bench.rs b/src/bench.rs
index caf4ed8..4abd5bf 100644
--- a/src/bench.rs
+++ b/src/bench.rs
@@ -1393,6 +1393,20 @@ mod tests {
         ));
         example(OPT);
     }
+
+    #[test]
+    fn example_rocksdb() {
+        let tmp_dir = tempfile::tempdir().unwrap();
+        let opt = format!(
+            r#"
+            [map]
+            name = "rocksdb"
+            path = "{}"
+            "#,
+            tmp_dir.path().to_str().unwrap().to_string()
+        );
+        example(&opt);
+    }
 }
 
 // }}} tests
diff --git a/src/server.rs b/src/server.rs
index 6946d6d..ad3c515 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -824,6 +824,28 @@ mod tests {
         simple(map);
     }
 
+    #[test]
+    fn simple_mutex_btreemap() {
+        let map = BenchKVMap::Regular(Box::new(btreemap::MutexBTreeMap::new()));
+        simple(map);
+    }
+
+    #[test]
+    fn simple_rwlock_btreemap() {
+        let map = BenchKVMap::Regular(Box::new(btreemap::RwLockBTreeMap::new()));
+        simple(map);
+    }
+
+    #[test]
+    fn simple_rocksdb() {
+        let tmp_dir = tempfile::tempdir().unwrap();
+        let opt = rocksdb::RocksDBOpt {
+            path: tmp_dir.path().to_str().unwrap().to_string(),
+        };
+        let map = BenchKVMap::Regular(Box::new(rocksdb::RocksDB::new(&opt)));
+        simple(map);
+    }
+
     fn batch(map: BenchKVMap) {
         const NR_CLIENTS: usize = 8;
         const NR_BATCHES: usize = 1000;
@@ -977,4 +999,14 @@ mod tests {
         let map = BenchKVMap::Regular(Box::new(btreemap::RwLockBTreeMap::new()));
         batch(map);
     }
+
+    #[test]
+    fn batch_rocksdb() {
+        let tmp_dir = tempfile::tempdir().unwrap();
+        let opt = rocksdb::RocksDBOpt {
+            path: tmp_dir.path().to_str().unwrap().to_string(),
+        };
+        let map = BenchKVMap::Regular(Box::new(rocksdb::RocksDB::new(&opt)));
+        batch(map);
+    }
 }
diff --git a/src/stores.rs b/src/stores.rs
index b67f739..a062a61 100644
--- a/src/stores.rs
+++ b/src/stores.rs
@@ -231,4 +231,15 @@ mod tests {
         let mut map = scc::SccHashMap::new();
         map_test(&mut map);
     }
+
+    #[test]
+    #[cfg(feature = "rocksdb")]
+    fn rocksdb() {
+        let tmp_dir = tempfile::tempdir().unwrap();
+        let opt = rocksdb::RocksDBOpt {
+            path: tmp_dir.path().to_str().unwrap().to_string(),
+        };
+        let mut map = rocksdb::RocksDB::new(&opt);
+        map_test(&mut map);
+    }
 }
diff --git a/src/stores/rocksdb.rs b/src/stores/rocksdb.rs
index 2d237f5..c658686 100644
--- a/src/stores/rocksdb.rs
+++ b/src/stores/rocksdb.rs
@@ -17,7 +17,7 @@ use serde::Deserialize;
 
 #[derive(Deserialize)]
 pub struct RocksDBOpt {
-    path: String,
+    pub path: String,
 }
 
 #[derive(Clone)]