[Storage] Fix tests. (aptos-labs#6970)

neoul · Mar 7, 2023 · eac8511 · eac8511
1 parent 1b1b627
commit eac8511
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/storage/aptosdb/src/pruner/state_store/test.rs b/storage/aptosdb/src/pruner/state_store/test.rs
@@ -378,7 +378,7 @@ fn verify_state_value_pruner(inputs: Vec<Vec<(StateKey, Option<StateValue>)>>) {
 
     let mut version = 0;
     let mut current_state_values = HashMap::new();
-    let pruner = StateKvPrunerManager::new(Arc::clone(&db.ledger_db), StateKvPrunerConfig {
+    let pruner = StateKvPrunerManager::new(Arc::clone(&db.state_kv_db), StateKvPrunerConfig {
         enable: true,
         prune_window: 0,
         batch_size: 1,

diff --git a/testsuite/smoke-test/src/genesis.rs b/testsuite/smoke-test/src/genesis.rs
@@ -257,7 +257,12 @@ async fn test_genesis_transaction_flow() {
 
     let db_dir = node.config().storage.dir();
     fs::remove_dir_all(&db_dir).unwrap();
-    db_restore(backup_path.path(), db_dir.as_path(), &[waypoint]);
+    db_restore(
+        backup_path.path(),
+        db_dir.as_path(),
+        &[waypoint],
+        node.config().storage.rocksdb_configs.use_state_kv_db,
+    );
 
     node.start().unwrap();
     wait_for_node(node, num_nodes - 2).await;

diff --git a/testsuite/smoke-test/src/storage.rs b/testsuite/smoke-test/src/storage.rs
@@ -128,7 +128,12 @@ async fn test_db_restore() {
 
     info!("---------- 3. stopped node 0, gonna restore DB.");
     // restore db from backup
-    db_restore(backup_path.path(), db_dir.as_path(), &[]);
+    db_restore(
+        backup_path.path(),
+        db_dir.as_path(),
+        &[],
+        node0_config.storage.rocksdb_configs.use_state_kv_db,
+    );
 
     expected_balance_0 -= 3;
     expected_balance_1 += 3;
@@ -370,7 +375,12 @@ pub(crate) fn db_backup(
     backup_path
 }
 
-pub(crate) fn db_restore(backup_path: &Path, db_path: &Path, trusted_waypoints: &[Waypoint]) {
+pub(crate) fn db_restore(
+    backup_path: &Path,
+    db_path: &Path,
+    trusted_waypoints: &[Waypoint],
+    use_state_kv_db: bool,
+) {
     let now = Instant::now();
     let bin_path = workspace_builder::get_bin("aptos-db-tool");
     let metadata_cache_path = TempPath::new();
@@ -384,6 +394,10 @@ pub(crate) fn db_restore(backup_path: &Path, db_path: &Path, trusted_waypoints:
         cmd.arg(&w.to_string());
     });
 
+    if use_state_kv_db {
+        cmd.arg("--use-state-kv-db");
+    }
+
     let status = cmd
         .args([
             "--target-db-dir",