deduplicate some test code (#3401)

kvstore/benches/basic.rs

@@ -4,11 +4,11 @@ extern crate test;
 use std::fs;
 use std::path::{Path, PathBuf};
 
-use rand::{self, thread_rng, Rng};
+use rand::{self, Rng};
 
 use test::Bencher;
 
-use solana_kvstore::{Config, Key, KvStore};
+use solana_kvstore::{test::gen, Config, Key, KvStore};
 
 const SMALL_SIZE: usize = 512;
 const LARGE_SIZE: usize = 32 * 1024;
@@ -44,7 +44,7 @@ fn bench_write_partitioned(bench: &mut Bencher, rows: &[(Key, Vec<u8>)], ledger_
 fn bench_write_small(bench: &mut Bencher) {
     let ledger_path = setup("bench_write_small");
     let num_entries = 32 * 1024;
-    let rows = gen_pairs(SMALL_SIZE).take(num_entries).collect::<Vec<_>>();
+    let rows = gen::pairs(SMALL_SIZE).take(num_entries).collect::<Vec<_>>();
     bench_write(bench, &rows, &ledger_path.to_string_lossy());
 }
 
@@ -53,7 +53,7 @@ fn bench_write_small(bench: &mut Bencher) {
 fn bench_write_small_partitioned(bench: &mut Bencher) {
     let ledger_path = setup("bench_write_small_partitioned");
     let num_entries = 32 * 1024;
-    let rows = gen_pairs(SMALL_SIZE).take(num_entries).collect::<Vec<_>>();
+    let rows = gen::pairs(SMALL_SIZE).take(num_entries).collect::<Vec<_>>();
     bench_write_partitioned(bench, &rows, &ledger_path.to_string_lossy());
 }
 
@@ -62,7 +62,7 @@ fn bench_write_small_partitioned(bench: &mut Bencher) {
 fn bench_write_large(bench: &mut Bencher) {
     let ledger_path = setup("bench_write_large");
     let num_entries = 32 * 1024;
-    let rows = gen_pairs(LARGE_SIZE).take(num_entries).collect::<Vec<_>>();
+    let rows = gen::pairs(LARGE_SIZE).take(num_entries).collect::<Vec<_>>();
     bench_write(bench, &rows, &ledger_path.to_string_lossy());
 }
 
@@ -71,7 +71,7 @@ fn bench_write_large(bench: &mut Bencher) {
 fn bench_write_huge(bench: &mut Bencher) {
     let ledger_path = setup("bench_write_huge");
     let num_entries = 32 * 1024;
-    let rows = gen_pairs(HUGE_SIZE).take(num_entries).collect::<Vec<_>>();
+    let rows = gen::pairs(HUGE_SIZE).take(num_entries).collect::<Vec<_>>();
     bench_write(bench, &rows, &ledger_path.to_string_lossy());
 }
 
@@ -86,8 +86,8 @@ fn bench_read_sequential(bench: &mut Bencher) {
     let num_large_blobs = 32 * 1024;
     let total_blobs = num_small_blobs + num_large_blobs;
 
-    let small = gen_data(SMALL_SIZE).take(num_small_blobs);
+    let small = gen::data(SMALL_SIZE).take(num_small_blobs);
-    let large = gen_data(LARGE_SIZE).take(num_large_blobs);
+    let large = gen::data(LARGE_SIZE).take(num_large_blobs);
     let rows = gen_seq_keys().zip(small.chain(large));
 
     let _ = store.put_many(rows);
@@ -119,8 +119,8 @@ fn bench_read_random(bench: &mut Bencher) {
     let num_large_blobs = 32 * 1024;
     let total_blobs = num_small_blobs + num_large_blobs;
 
-    let small = gen_data(SMALL_SIZE).take(num_small_blobs);
+    let small = gen::data(SMALL_SIZE).take(num_small_blobs);
-    let large = gen_data(LARGE_SIZE).take(num_large_blobs);
+    let large = gen::data(LARGE_SIZE).take(num_large_blobs);
     let rows = gen_seq_keys().zip(small.chain(large));
 
     let _ = store.put_many(rows);
@@ -165,24 +165,6 @@ fn gen_seq_keys() -> impl Iterator<Item = Key> {
     })
 }
 
-fn gen_keys() -> impl Iterator<Item = Key> {
-    let mut rng = thread_rng();
-
-    std::iter::repeat_with(move || {
-        let buf = rng.gen();
-
-        Key(buf)
-    })
-}
-
-fn gen_data(size: usize) -> impl Iterator<Item = Vec<u8>> {
-    std::iter::repeat(vec![1u8; size])
-}
-
-fn gen_pairs(data_size: usize) -> impl Iterator<Item = (Key, Vec<u8>)> {
-    gen_keys().zip(gen_data(data_size))
-}
-
 fn teardown<P: AsRef<Path>>(p: P) {
     KvStore::destroy(p).expect("Expect successful store destruction");
 }

kvstore/src/lib.rs

@@ -368,3 +368,41 @@ fn is_lvl0_full(tables: &[BTreeMap<Key, SSTable>], config: &Config) -> bool {
         tables[0].len() > config.max_tables
     }
 }
+
+#[cfg(any(test, debug_assertions))]
+pub mod test {
+    pub mod gen {
+        use crate::Key;
+        use rand::distributions::Uniform;
+        use rand::{rngs::SmallRng, FromEntropy, Rng};
+        use std::iter;
+        use std::ops::Range;
+
+        pub fn keys() -> impl Iterator<Item = Key> {
+            let mut rng = SmallRng::from_entropy();
+            iter::repeat_with(move || Key(rng.gen()))
+        }
+
+        pub fn data(size: usize) -> impl Iterator<Item = Vec<u8>> {
+            iter::repeat(vec![0; size])
+        }
+
+        pub fn data_vary(range: Range<u64>) -> impl Iterator<Item = Vec<u8>> {
+            let dist = Uniform::from(range);
+            let mut rng = SmallRng::from_entropy();
+
+            iter::repeat_with(move || {
+                let size: u64 = rng.sample(dist);
+                vec![0; size as usize]
+            })
+        }
+
+        pub fn pairs(size: usize) -> impl Iterator<Item = (Key, Vec<u8>)> {
+            keys().zip(data(size))
+        }
+
+        pub fn pairs_vary(range: Range<u64>) -> impl Iterator<Item = (Key, Vec<u8>)> {
+            keys().zip(data_vary(range))
+        }
+    }
+}

kvstore/src/sstable.rs

@@ -410,9 +410,9 @@ fn overlapping<T: Ord + Eq>(r1: &RangeInclusive<T>, r2: &RangeInclusive<T>) -> b
 }
 
 #[cfg(test)]
-mod test {
+pub mod test {
     use super::*;
-    use rand::{thread_rng, Rng};
+    use crate::test::gen;
     use std::sync::{Arc, RwLock};
 
     #[test]
@@ -557,17 +557,8 @@ mod test {
     }
 
     fn gen_records() -> impl Iterator<Item = (Key, Value)> {
-        let mut rng = thread_rng();
+        gen::pairs_vary(0..255)
-        let commit = rng.gen();
+            .map(|(key, bytes)| (key, Value::new(bytes.len() as i64, Some(bytes))))
-
-        std::iter::repeat_with(move || {
-            let buf: [u8; KEY_LEN] = rng.gen();
-            let data_size: u8 = buf[0];
-
-            let val = Some(vec![0; data_size as usize]);
-
-            (Key(buf), Value::new(commit, val))
-        })
     }
 
 }

kvstore/src/storage.rs

@@ -158,7 +158,7 @@ fn opt_bytes_memory(bytes: &Option<Vec<u8>>) -> usize {
 #[cfg(test)]
 mod test {
     use super::*;
-    use rand::{self, thread_rng, Rng};
+    use crate::test::gen;
 
     const COMMIT: i64 = -1;
 
@@ -168,7 +168,7 @@ mod test {
 
         let mut table = MemTable::default();
 
-        for (key, data) in gen_pairs(DATA_SIZE).take(1024) {
+        for (key, data) in gen::pairs(DATA_SIZE).take(1024) {
             table.put(&key, COMMIT, &data);
         }
 
@@ -181,7 +181,7 @@ mod test {
         const DATA_SIZE: usize = 32;
 
         let mut table = MemTable::default();
-        let input = gen_pairs(DATA_SIZE).take(1024).collect::<Vec<_>>();
+        let input = gen::pairs(DATA_SIZE).take(1024).collect::<Vec<_>>();
 
         for (key, data) in &input {
             table.put(key, COMMIT, data);
@@ -196,7 +196,7 @@ mod test {
         assert_eq!(table.mem_size, expected_size);
 
         // Deletes of things not in the memory table must be recorded
-        for key in gen_keys().take(512) {
+        for key in gen::keys().take(512) {
             table.delete(&key, COMMIT);
         }
 
@@ -207,8 +207,8 @@ mod test {
     #[test]
     fn test_put_order_irrelevant() {
         let (mut table_1, mut table_2) = (MemTable::default(), MemTable::default());
-        let big_input: Vec<_> = gen_pairs(1024).take(128).collect();
+        let big_input: Vec<_> = gen::pairs(1024).take(128).collect();
-        let small_input: Vec<_> = gen_pairs(16).take(128).collect();
+        let small_input: Vec<_> = gen::pairs(16).take(128).collect();
 
         for (key, data) in big_input.iter().chain(small_input.iter()) {
             table_1.put(key, COMMIT, data);
@@ -237,8 +237,8 @@ mod test {
     #[test]
     fn test_delete_order_irrelevant() {
         let (mut table_1, mut table_2) = (MemTable::default(), MemTable::default());
-        let big_input: Vec<_> = gen_pairs(1024).take(128).collect();
+        let big_input: Vec<_> = gen::pairs(1024).take(128).collect();
-        let small_input: Vec<_> = gen_pairs(16).take(128).collect();
+        let small_input: Vec<_> = gen::pairs(16).take(128).collect();
 
         for (key, data) in big_input.iter().chain(small_input.iter()) {
             table_1.put(key, COMMIT, data);
@@ -277,22 +277,4 @@ mod test {
         assert_eq!(table_1.mem_size, table_2.mem_size);
         assert_eq!(table_1.values, table_2.values);
     }
-
-    fn gen_keys() -> impl Iterator<Item = Key> {
-        let mut rng = thread_rng();
-
-        std::iter::repeat_with(move || {
-            let buf = rng.gen();
-
-            Key(buf)
-        })
-    }
-
-    fn gen_data(size: usize) -> impl Iterator<Item = Vec<u8>> {
-        std::iter::repeat(vec![1u8; size])
-    }
-
-    fn gen_pairs(data_size: usize) -> impl Iterator<Item = (Key, Vec<u8>)> {
-        gen_keys().zip(gen_data(data_size))
-    }
 }

kvstore/src/writebatch.rs

@@ -118,15 +118,15 @@ impl Default for Config {
 #[cfg(test)]
 mod test {
     use super::*;
+    use crate::test::gen;
     use crate::writelog::Config as WalConfig;
-    use rand::{self, thread_rng, Rng};
 
     const CAPACITY: usize = 10 * 1024;
 
     #[test]
     fn test_put_associative() {
         let mut writebatch = setup();
-        let input: Vec<_> = gen_pairs(32).take(100).collect();
+        let input: Vec<_> = gen::pairs(32).take(100).collect();
 
         writebatch.put_many(input.iter()).unwrap();
 
@@ -146,7 +146,7 @@ mod test {
     #[test]
     fn test_delete_associative() {
         let (mut writebatch, mut writebatch2) = (setup(), setup());
-        let input: Vec<_> = gen_pairs(32).take(100).collect();
+        let input: Vec<_> = gen::pairs(32).take(100).collect();
 
         writebatch.put_many(input.iter()).unwrap();
         writebatch2.put_many(input.iter()).unwrap();
@@ -172,7 +172,7 @@ mod test {
         let mut writebatch = setup();
 
         let space_per_record = CAPACITY / AMT_RECORDS - MemTable::OVERHEAD_PER_RECORD;
-        let input: Vec<_> = gen_pairs(space_per_record).take(AMT_RECORDS).collect();
+        let input: Vec<_> = gen::pairs(space_per_record).take(AMT_RECORDS).collect();
 
         writebatch.put_many(input.iter()).unwrap();
 
@@ -181,7 +181,7 @@ mod test {
             _ => panic!("Writebatch should be exactly at capacity"),
         }
 
-        let (key, data) = gen_pairs(space_per_record).next().unwrap();
+        let (key, data) = gen::pairs(space_per_record).next().unwrap();
         let result = writebatch.put(&key, &data);
         assert!(result.is_err());
 
@@ -206,22 +206,4 @@ mod test {
             log: Arc::new(RwLock::new(log)),
         }
     }
-
-    fn gen_keys() -> impl Iterator<Item = Key> {
-        let mut rng = thread_rng();
-
-        std::iter::repeat_with(move || {
-            let buf = rng.gen();
-
-            Key(buf)
-        })
-    }
-
-    fn gen_data(size: usize) -> impl Iterator<Item = Vec<u8>> {
-        std::iter::repeat(vec![1u8; size])
-    }
-
-    fn gen_pairs(data_size: usize) -> impl Iterator<Item = (Key, Vec<u8>)> {
-        gen_keys().zip(gen_data(data_size))
-    }
 }

kvstore/tests/basic.rs

@@ -1,8 +1,7 @@
-use rand::{thread_rng, Rng};
-
 use std::fs;
 use std::path::{Path, PathBuf};
 
+use solana_kvstore::test::gen;
 use solana_kvstore::{Config, Key, KvStore};
 
 const KB: usize = 1024;
@@ -20,7 +19,7 @@ fn test_put_get() {
     };
 
     let lsm = KvStore::open(&path, cfg).unwrap();
-    let (key, bytes) = gen_pairs(HALF_KB).take(1).next().unwrap();
+    let (key, bytes) = gen::pairs(HALF_KB).take(1).next().unwrap();
 
     lsm.put(&key, &bytes).expect("put fail");
     let out_bytes = lsm.get(&key).expect("get fail").expect("missing");
@@ -42,7 +41,7 @@ fn test_put_get_many() {
     };
     let lsm = KvStore::open(&path, cfg).unwrap();
 
-    let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(1024).collect();
+    let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(1024).collect();
     pairs.sort_unstable_by_key(|(k, _)| *k);
 
     lsm.put_many(pairs.clone().drain(..))
@@ -70,7 +69,7 @@ fn test_delete() {
     };
     let lsm = KvStore::open(&path, cfg).unwrap();
 
-    let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(64 * 6).collect();
+    let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(64 * 6).collect();
     pairs.sort_unstable_by_key(|(k, _)| *k);
 
     for (k, i) in pairs.iter() {
@@ -104,7 +103,7 @@ fn test_delete_many() {
     };
     let lsm = KvStore::open(&path, cfg).unwrap();
 
-    let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(64 * 6).collect();
+    let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(64 * 6).collect();
     pairs.sort_unstable_by_key(|(k, _)| *k);
 
     for (k, i) in pairs.iter() {
@@ -132,7 +131,7 @@ fn test_close_reopen() {
     let cfg = Config::default();
     let lsm = KvStore::open(&path, cfg).unwrap();
 
-    let mut pairs: Vec<_> = gen_pairs(KB).take(1024).collect();
+    let mut pairs: Vec<_> = gen::pairs(KB).take(1024).collect();
     pairs.sort_unstable_by_key(|(k, _)| *k);
 
     for (k, i) in pairs.iter() {
@@ -174,7 +173,7 @@ fn test_partitioned() {
 
     let lsm = KvStore::partitioned(&path, &storage_dirs, cfg).unwrap();
 
-    let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(64 * 12).collect();
+    let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(64 * 12).collect();
     pairs.sort_unstable_by_key(|(k, _)| *k);
 
     lsm.put_many(pairs.iter()).expect("put_many fail");
@@ -207,7 +206,7 @@ fn test_in_memory() {
     };
     let lsm = KvStore::open(&path, cfg).unwrap();
 
-    let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(64 * 12).collect();
+    let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(64 * 12).collect();
     pairs.sort_unstable_by_key(|(k, _)| *k);
 
     lsm.put_many(pairs.iter()).expect("put_many fail");
@@ -239,14 +238,3 @@ fn setup(test_name: &str) -> PathBuf {
 fn teardown(p: &Path) {
     KvStore::destroy(p).expect("Expect successful store destruction");
 }
-
-fn gen_pairs(data_size: usize) -> impl Iterator<Item = (Key, Vec<u8>)> {
-    let mut rng = thread_rng();
-
-    std::iter::repeat_with(move || {
-        let data = vec![0u8; data_size];
-        let buf = rng.gen();
-
-        (Key(buf), data)
-    })
-}