From 4b349f6d47c0cbd08457e6e383fa68066678a7b4 Mon Sep 17 00:00:00 2001
From: steviez <steven@anza.xyz>
Date: Mon, 9 Dec 2024 23:19:03 -0600
Subject: [PATCH] More test cleanup

---
 ledger/src/blockstore/blockstore_purge.rs | 89 ++++++++---------------
 ledger/src/blockstore_db.rs               | 22 ++----
 2 files changed, 37 insertions(+), 74 deletions(-)

diff --git a/ledger/src/blockstore/blockstore_purge.rs b/ledger/src/blockstore/blockstore_purge.rs
index db9dc001d99021..66cf51e7970e63 100644
--- a/ledger/src/blockstore/blockstore_purge.rs
+++ b/ledger/src/blockstore/blockstore_purge.rs
@@ -804,18 +804,17 @@ pub mod tests {
     }
 
     fn get_index_bounds(blockstore: &Blockstore) -> (Box<[u8]>, Box<[u8]>) {
-        let first_index = {
-            let mut status_entry_iterator = blockstore
-                .transaction_status_cf
-                .iterator_cf_raw_key(IteratorMode::Start);
-            status_entry_iterator.next().unwrap().unwrap().0
-        };
-        let last_index = {
-            let mut status_entry_iterator = blockstore
-                .transaction_status_cf
-                .iterator_cf_raw_key(IteratorMode::End);
-            status_entry_iterator.next().unwrap().unwrap().0
-        };
+        let (first_index, _value) = blockstore
+            .transaction_status_cf
+            .iterator_cf_raw_key(IteratorMode::Start)
+            .next()
+            .unwrap();
+        let (last_index, _value) = blockstore
+            .transaction_status_cf
+            .iterator_cf_raw_key(IteratorMode::End)
+            .next()
+            .unwrap();
+
         (first_index, last_index)
     }
 
@@ -990,27 +989,13 @@ pub mod tests {
         let max_slot = 19;
 
         clear_and_repopulate_transaction_statuses_for_test(&blockstore, max_slot);
-        let first_index = {
-            let mut status_entry_iterator = blockstore
-                .transaction_status_cf
-                .iter(IteratorMode::Start)
-                .unwrap();
-            status_entry_iterator.next().unwrap().0
-        };
-        let last_index = {
-            let mut status_entry_iterator = blockstore
-                .transaction_status_cf
-                .iter(IteratorMode::End)
-                .unwrap();
-            status_entry_iterator.next().unwrap().0
-        };
+        let (first_index, last_index) = get_index_bounds(&blockstore);
 
         let oldest_slot = 3;
         blockstore.db.set_oldest_slot(oldest_slot);
-        blockstore.transaction_status_cf.compact_range_raw_key(
-            &cf::TransactionStatus::key(first_index),
-            &cf::TransactionStatus::key(last_index),
-        );
+        blockstore
+            .transaction_status_cf
+            .compact_range_raw_key(&first_index, &last_index);
 
         let status_entry_iterator = blockstore
             .transaction_status_cf
@@ -1024,27 +1009,13 @@ pub mod tests {
         assert_eq!(count, max_slot - (oldest_slot - 1));
 
         clear_and_repopulate_transaction_statuses_for_test(&blockstore, max_slot);
-        let first_index = {
-            let mut status_entry_iterator = blockstore
-                .transaction_status_cf
-                .iter(IteratorMode::Start)
-                .unwrap();
-            status_entry_iterator.next().unwrap().0
-        };
-        let last_index = {
-            let mut status_entry_iterator = blockstore
-                .transaction_status_cf
-                .iter(IteratorMode::End)
-                .unwrap();
-            status_entry_iterator.next().unwrap().0
-        };
+        let (first_index, last_index) = get_index_bounds(&blockstore);
 
         let oldest_slot = 12;
         blockstore.db.set_oldest_slot(oldest_slot);
-        blockstore.transaction_status_cf.compact_range_raw_key(
-            &cf::TransactionStatus::key(first_index),
-            &cf::TransactionStatus::key(last_index),
-        );
+        blockstore
+            .transaction_status_cf
+            .compact_range_raw_key(&first_index, &last_index);
 
         let status_entry_iterator = blockstore
             .transaction_status_cf
@@ -1100,18 +1071,16 @@ pub mod tests {
             )
             .unwrap();
 
-        let first_index = {
-            let mut memos_iterator = blockstore
-                .transaction_memos_cf
-                .iterator_cf_raw_key(IteratorMode::Start);
-            memos_iterator.next().unwrap().unwrap().0
-        };
-        let last_index = {
-            let mut memos_iterator = blockstore
-                .transaction_memos_cf
-                .iterator_cf_raw_key(IteratorMode::End);
-            memos_iterator.next().unwrap().unwrap().0
-        };
+        let (first_index, _value) = blockstore
+            .transaction_memos_cf
+            .iterator_cf_raw_key(IteratorMode::Start)
+            .next()
+            .unwrap();
+        let (last_index, _value) = blockstore
+            .transaction_memos_cf
+            .iterator_cf_raw_key(IteratorMode::End)
+            .next()
+            .unwrap();
 
         // Purge at slot 0 should not affect any memos
         blockstore.db.set_oldest_slot(0);
diff --git a/ledger/src/blockstore_db.rs b/ledger/src/blockstore_db.rs
index b0e520ddb2c65e..14f31163c2067b 100644
--- a/ledger/src/blockstore_db.rs
+++ b/ledger/src/blockstore_db.rs
@@ -2197,20 +2197,14 @@ pub mod tests {
     {
         pub(crate) fn iterator_cf_raw_key(
             &self,
-            iterator_mode: IteratorMode<Vec<u8>>,
-        ) -> DBIterator {
-            let cf = self.handle();
-            let start_key;
-            let iterator_mode = match iterator_mode {
-                IteratorMode::Start => RocksIteratorMode::Start,
-                IteratorMode::End => RocksIteratorMode::End,
-                IteratorMode::From(start_from, direction) => {
-                    start_key = start_from.clone();
-                    RocksIteratorMode::From(&start_key, direction)
-                }
-            };
-
-            self.backend.iterator_cf(cf, iterator_mode)
+            iterator_mode: IteratorMode<C::Index>,
+        ) -> impl Iterator<Item = (Box<[u8]>, Box<[u8]>)> + '_ {
+            // The conversion of key back into Box<[u8]> incurs an extra
+            // allocation. However, this is test code and the goal is to
+            // maximize code reuse over efficiency
+            self.iter(iterator_mode)
+                .unwrap()
+                .map(|(key, value)| (Box::from(C::key(key)), value))
         }
     }
 }