deduplicate some test code (#3401)

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();
-        let commit = rng.gen();
-
-        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))
-        })
+        gen::pairs_vary(0..255)
+            .map(|(key, bytes)| (key, Value::new(bytes.len() as i64, Some(bytes))))
     }
 
 }

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 small_input: Vec<_> = gen_pairs(16).take(128).collect();
+        let big_input: Vec<_> = gen::pairs(1024).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 small_input: Vec<_> = gen_pairs(16).take(128).collect();
+        let big_input: Vec<_> = gen::pairs(1024).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))
-    }
 }