From 908049d9e7d130cfeee31d144591e77b69b461d3 Mon Sep 17 00:00:00 2001 From: Michal Nazarewicz Date: Tue, 12 Jul 2022 09:53:29 +0200 Subject: [PATCH] Refactor iterator-related tests (#652) --- tests/test_db.rs | 36 +++--- tests/test_iterator.rs | 235 +++++++++++----------------------- tests/test_slice_transform.rs | 39 ++---- tests/util/mod.rs | 27 +++- 4 files changed, 129 insertions(+), 208 deletions(-) diff --git a/tests/test_db.rs b/tests/test_db.rs index 6ceed76..3a70068 100644 --- a/tests/test_db.rs +++ b/tests/test_db.rs @@ -25,7 +25,7 @@ use rocksdb::{ PerfContext, PerfMetric, ReadOptions, SingleThreaded, SliceTransform, Snapshot, UniversalCompactOptions, UniversalCompactionStopStyle, WriteBatch, DB, }; -use util::DBPath; +use util::{assert_iter, pair, DBPath}; #[test] fn external() { @@ -213,12 +213,10 @@ fn iterator_test_upper_bound() { let mut readopts = ReadOptions::default(); readopts.set_iterate_upper_bound(b"k4".to_vec()); - let iter = db.iterator_opt(IteratorMode::Start, readopts); - let expected: Vec<_> = vec![(b"k1", b"v1"), (b"k2", b"v2"), (b"k3", b"v3")] - .into_iter() - .map(|(k, v)| (k.to_vec().into_boxed_slice(), v.to_vec().into_boxed_slice())) - .collect(); - assert_eq!(expected, iter.collect::>()); + assert_iter( + db.iterator_opt(IteratorMode::Start, readopts), + &[pair(b"k1", b"v1"), pair(b"k2", b"v2"), pair(b"k3", b"v3")], + ); } } @@ -236,12 +234,10 @@ fn iterator_test_lower_bound() { let mut readopts = ReadOptions::default(); readopts.set_iterate_lower_bound(b"k4".to_vec()); - let iter = db.iterator_opt(IteratorMode::Start, readopts); - let expected: Vec<_> = vec![(b"k4", b"v4"), (b"k5", b"v5")] - .into_iter() - .map(|(k, v)| (k.to_vec().into_boxed_slice(), v.to_vec().into_boxed_slice())) - .collect(); - assert_eq!(expected, iter.collect::>()); + assert_iter( + db.iterator_opt(IteratorMode::Start, readopts), + &[pair(b"k4", b"v4"), pair(b"k5", b"v5")], + ); } } @@ -784,12 +780,14 @@ fn prefix_extract_and_iterate_test() { readopts.set_iterate_lower_bound(b"p1".to_vec()); readopts.set_pin_data(true); - let iter = db.iterator_opt(IteratorMode::Start, readopts); - let expected: Vec<_> = vec![(b"p1_k1", b"v1"), (b"p1_k3", b"v3"), (b"p1_k4", b"v4")] - .into_iter() - .map(|(k, v)| (k.to_vec().into_boxed_slice(), v.to_vec().into_boxed_slice())) - .collect(); - assert_eq!(expected, iter.collect::>()); + assert_iter( + db.iterator_opt(IteratorMode::Start, readopts), + &[ + pair(b"p1_k1", b"v1"), + pair(b"p1_k3", b"v3"), + pair(b"p1_k4", b"v4"), + ], + ); } } diff --git a/tests/test_iterator.rs b/tests/test_iterator.rs index 66215a2..5f561fa 100644 --- a/tests/test_iterator.rs +++ b/tests/test_iterator.rs @@ -17,130 +17,61 @@ mod util; use pretty_assertions::assert_eq; use rocksdb::{Direction, IteratorMode, MemtableFactory, Options, DB}; -use util::DBPath; - -fn cba(input: &[u8]) -> Box<[u8]> { - input.to_vec().into_boxed_slice() -} +use util::{assert_iter, assert_iter_reversed, pair, DBPath}; #[test] #[allow(clippy::cognitive_complexity)] fn test_iterator() { let n = DBPath::new("_rust_rocksdb_iterator_test"); { - let k1: Box<[u8]> = b"k1".to_vec().into_boxed_slice(); - let k2: Box<[u8]> = b"k2".to_vec().into_boxed_slice(); - let k3: Box<[u8]> = b"k3".to_vec().into_boxed_slice(); - let k4: Box<[u8]> = b"k4".to_vec().into_boxed_slice(); - let v1: Box<[u8]> = b"v1111".to_vec().into_boxed_slice(); - let v2: Box<[u8]> = b"v2222".to_vec().into_boxed_slice(); - let v3: Box<[u8]> = b"v3333".to_vec().into_boxed_slice(); - let v4: Box<[u8]> = b"v4444".to_vec().into_boxed_slice(); + const K1: &[u8] = b"k1"; + const K2: &[u8] = b"k2"; + const K3: &[u8] = b"k3"; + const K4: &[u8] = b"k4"; + const V1: &[u8] = b"v1111"; + const V2: &[u8] = b"v2222"; + const V3: &[u8] = b"v3333"; + const V4: &[u8] = b"v4444"; + let db = DB::open_default(&n).unwrap(); - let p = db.put(&*k1, &*v1); - assert!(p.is_ok()); - let p = db.put(&*k2, &*v2); - assert!(p.is_ok()); - let p = db.put(&*k3, &*v3); - assert!(p.is_ok()); - let expected = vec![ - (cba(&k1), cba(&v1)), - (cba(&k2), cba(&v2)), - (cba(&k3), cba(&v3)), - ]; - { - let iterator1 = db.iterator(IteratorMode::Start); - assert_eq!(iterator1.collect::>(), expected); - } + assert!(db.put(K1, V1).is_ok()); + assert!(db.put(K2, V2).is_ok()); + assert!(db.put(K3, V3).is_ok()); + let expected = [pair(K1, V1), pair(K2, V2), pair(K3, V3)]; + assert_iter(db.iterator(IteratorMode::Start), &expected); // Test that it's idempotent - { - let iterator1 = db.iterator(IteratorMode::Start); - assert_eq!(iterator1.collect::>(), expected); - } - { - let iterator1 = db.iterator(IteratorMode::Start); - assert_eq!(iterator1.collect::>(), expected); - } - { - let iterator1 = db.iterator(IteratorMode::Start); - assert_eq!(iterator1.collect::>(), expected); - } + assert_iter(db.iterator(IteratorMode::Start), &expected); + assert_iter(db.iterator(IteratorMode::Start), &expected); + assert_iter(db.iterator(IteratorMode::Start), &expected); // Test it in reverse a few times - { - let iterator1 = db.iterator(IteratorMode::End); - let mut tmp_vec = iterator1.collect::>(); - tmp_vec.reverse(); - assert_eq!(tmp_vec, expected); - } - { - let iterator1 = db.iterator(IteratorMode::End); - let mut tmp_vec = iterator1.collect::>(); - tmp_vec.reverse(); - assert_eq!(tmp_vec, expected); - } - { - let iterator1 = db.iterator(IteratorMode::End); - let mut tmp_vec = iterator1.collect::>(); - tmp_vec.reverse(); - assert_eq!(tmp_vec, expected); - } - { - let iterator1 = db.iterator(IteratorMode::End); - let mut tmp_vec = iterator1.collect::>(); - tmp_vec.reverse(); - assert_eq!(tmp_vec, expected); - } - { - let iterator1 = db.iterator(IteratorMode::End); - let mut tmp_vec = iterator1.collect::>(); - tmp_vec.reverse(); - assert_eq!(tmp_vec, expected); - } + assert_iter_reversed(db.iterator(IteratorMode::End), &expected); + assert_iter_reversed(db.iterator(IteratorMode::End), &expected); + assert_iter_reversed(db.iterator(IteratorMode::End), &expected); + assert_iter_reversed(db.iterator(IteratorMode::End), &expected); // Try it forward again - { - let iterator1 = db.iterator(IteratorMode::Start); - assert_eq!(iterator1.collect::>(), expected); - } - { - let iterator1 = db.iterator(IteratorMode::Start); - assert_eq!(iterator1.collect::>(), expected); - } + assert_iter(db.iterator(IteratorMode::Start), &expected); + assert_iter(db.iterator(IteratorMode::Start), &expected); + + { + let old_iterator = db.iterator(IteratorMode::Start); + assert!(db.put(K4, V4).is_ok()); + assert_iter(old_iterator, &expected); + } + let expected2 = [pair(K1, V1), pair(K2, V2), pair(K3, V3), pair(K4, V4)]; + assert_iter(db.iterator(IteratorMode::Start), &expected2); + assert_iter( + db.iterator(IteratorMode::From(b"k2", Direction::Forward)), + &expected2[1..], + ); + assert_iter_reversed( + db.iterator(IteratorMode::From(b"k2", Direction::Reverse)), + &expected[..2], + ); + assert_iter_reversed( + db.iterator(IteratorMode::From(b"zz", Direction::Reverse)), + &expected2, + ); - let old_iterator = db.iterator(IteratorMode::Start); - let p = db.put(&*k4, &*v4); - assert!(p.is_ok()); - let expected2 = vec![ - (cba(&k1), cba(&v1)), - (cba(&k2), cba(&v2)), - (cba(&k3), cba(&v3)), - (cba(&k4), cba(&v4)), - ]; - { - assert_eq!(old_iterator.collect::>(), expected); - } - { - let iterator1 = db.iterator(IteratorMode::Start); - assert_eq!(iterator1.collect::>(), expected2); - } - { - let iterator1 = db.iterator(IteratorMode::From(b"k2", Direction::Forward)); - let expected = vec![ - (cba(&k2), cba(&v2)), - (cba(&k3), cba(&v3)), - (cba(&k4), cba(&v4)), - ]; - assert_eq!(iterator1.collect::>(), expected); - } - { - let iterator1 = db.iterator(IteratorMode::From(b"k2", Direction::Reverse)); - let expected = vec![(cba(&k2), cba(&v2)), (cba(&k1), cba(&v1))]; - assert_eq!(iterator1.collect::>(), expected); - } - { - let iterator1 = db.iterator(IteratorMode::From(b"zz", Direction::Reverse)); - let expected = vec![(cba(&k4), cba(&v4)), (cba(&k3), cba(&v3))]; - assert_eq!(iterator1.take(2).collect::>(), expected); - } { let iterator1 = db.iterator(IteratorMode::From(b"k0", Direction::Forward)); assert!(iterator1.valid()); @@ -169,18 +100,14 @@ fn test_iterator() { } } -fn key(k: &[u8]) -> Box<[u8]> { - k.to_vec().into_boxed_slice() -} - #[test] fn test_prefix_iterator() { let n = DBPath::new("_rust_rocksdb_prefix_iterator_test"); { - let a1: Box<[u8]> = key(b"aaa1"); - let a2: Box<[u8]> = key(b"aaa2"); - let b1: Box<[u8]> = key(b"bbb1"); - let b2: Box<[u8]> = key(b"bbb2"); + const A1: &[u8] = b"aaa1"; + const A2: &[u8] = b"aaa2"; + const B1: &[u8] = b"bbb1"; + const B2: &[u8] = b"bbb2"; let prefix_extractor = rocksdb::SliceTransform::create_fixed_prefix(3); @@ -190,22 +117,14 @@ fn test_prefix_iterator() { let db = DB::open(&opts, &n).unwrap(); - assert!(db.put(&*a1, &*a1).is_ok()); - assert!(db.put(&*a2, &*a2).is_ok()); - assert!(db.put(&*b1, &*b1).is_ok()); - assert!(db.put(&*b2, &*b2).is_ok()); - - { - let expected = vec![(cba(&a1), cba(&a1)), (cba(&a2), cba(&a2))]; - let a_iterator = db.prefix_iterator(b"aaa"); - assert_eq!(a_iterator.collect::>(), expected) - } + assert!(db.put(A1, A1).is_ok()); + assert!(db.put(A2, A2).is_ok()); + assert!(db.put(B1, B1).is_ok()); + assert!(db.put(B2, B2).is_ok()); - { - let expected = vec![(cba(&b1), cba(&b1)), (cba(&b2), cba(&b2))]; - let b_iterator = db.prefix_iterator(b"bbb"); - assert_eq!(b_iterator.collect::>(), expected) - } + assert_iter(db.prefix_iterator(b"aaa"), &[pair(A1, A1), pair(A2, A2)]); + assert_iter(db.prefix_iterator(b"bbb"), &[pair(B1, B1), pair(B2, B2)]); + assert_iter(db.prefix_iterator(A2), &[pair(A2, A2)]); } } @@ -244,13 +163,14 @@ fn test_prefix_iterator_uses_full_prefix() { assert!(db.put(key, *value).is_ok()); } - let prefix = [0, 1, 1]; - let results: Vec<_> = db - .prefix_iterator(&prefix) - .map(|(_, v)| std::str::from_utf8(&v).unwrap().to_string()) - .collect(); - - assert_eq!(results, vec!("444", "555", "666")); + assert_iter( + db.prefix_iterator(&[0, 1, 1]), + &[ + pair(&[0, 1, 1, 1], b"444"), + pair(&[0, 1, 2, 1], b"555"), + pair(&[0, 2, 0, 0], b"666"), + ], + ); } } @@ -258,10 +178,10 @@ fn test_prefix_iterator_uses_full_prefix() { fn test_full_iterator() { let path = DBPath::new("full_iterator_test"); { - let a1: Box<[u8]> = key(b"aaa1"); - let a2: Box<[u8]> = key(b"aaa2"); - let b1: Box<[u8]> = key(b"bbb1"); - let b2: Box<[u8]> = key(b"bbb2"); + const A1: &[u8] = b"aaa1"; + const A2: &[u8] = b"aaa2"; + const B1: &[u8] = b"bbb1"; + const B2: &[u8] = b"bbb2"; let prefix_extractor = rocksdb::SliceTransform::create_fixed_prefix(3); let factory = MemtableFactory::HashSkipList { @@ -278,25 +198,20 @@ fn test_full_iterator() { let db = DB::open(&opts, &path).unwrap(); - assert!(db.put(&*a1, &*a1).is_ok()); - assert!(db.put(&*a2, &*a2).is_ok()); - assert!(db.put(&*b1, &*b1).is_ok()); - assert!(db.put(&*b2, &*b2).is_ok()); + assert!(db.put(A1, A1).is_ok()); + assert!(db.put(A2, A2).is_ok()); + assert!(db.put(B1, B1).is_ok()); + assert!(db.put(B2, B2).is_ok()); // A normal iterator won't work here since we're using a HashSkipList for our memory table // implementation (which buckets keys based on their prefix): let bad_iterator = db.iterator(IteratorMode::Start); assert_eq!(bad_iterator.collect::>(), vec![]); - let expected = vec![ - (cba(&a1), cba(&a1)), - (cba(&a2), cba(&a2)), - (cba(&b1), cba(&b1)), - (cba(&b2), cba(&b2)), - ]; - - let a_iterator = db.full_iterator(IteratorMode::Start); - assert_eq!(a_iterator.collect::>(), expected) + assert_iter( + db.full_iterator(IteratorMode::Start), + &[pair(A1, A1), pair(A2, A2), pair(B1, B1), pair(B2, B2)], + ); } } diff --git a/tests/test_slice_transform.rs b/tests/test_slice_transform.rs index c2633d8..afbb248 100644 --- a/tests/test_slice_transform.rs +++ b/tests/test_slice_transform.rs @@ -17,16 +17,16 @@ mod util; use pretty_assertions::assert_eq; use rocksdb::{Options, SliceTransform, DB}; -use util::DBPath; +use util::{assert_iter, pair, DBPath}; #[test] pub fn test_slice_transform() { let db_path = DBPath::new("_rust_rocksdb_slice_transform_test"); { - let a1: Box<[u8]> = key(b"aaa1"); - let a2: Box<[u8]> = key(b"aaa2"); - let b1: Box<[u8]> = key(b"bbb1"); - let b2: Box<[u8]> = key(b"bbb2"); + const A1: &[u8] = b"aaa1"; + const A2: &[u8] = b"aaa2"; + const B1: &[u8] = b"bbb1"; + const B2: &[u8] = b"bbb2"; fn first_three(k: &[u8]) -> &[u8] { &k[..3] @@ -40,30 +40,13 @@ pub fn test_slice_transform() { let db = DB::open(&opts, &db_path).unwrap(); - assert!(db.put(&*a1, &*a1).is_ok()); - assert!(db.put(&*a2, &*a2).is_ok()); - assert!(db.put(&*b1, &*b1).is_ok()); - assert!(db.put(&*b2, &*b2).is_ok()); + assert!(db.put(A1, A1).is_ok()); + assert!(db.put(A2, A2).is_ok()); + assert!(db.put(B1, B1).is_ok()); + assert!(db.put(B2, B2).is_ok()); - fn cba(input: &[u8]) -> Box<[u8]> { - input.to_vec().into_boxed_slice() - } - - fn key(k: &[u8]) -> Box<[u8]> { - k.to_vec().into_boxed_slice() - } - - { - let expected = vec![(cba(&a1), cba(&a1)), (cba(&a2), cba(&a2))]; - let a_iterator = db.prefix_iterator(b"aaa"); - assert_eq!(a_iterator.collect::>(), expected) - } - - { - let expected = vec![(cba(&b1), cba(&b1)), (cba(&b2), cba(&b2))]; - let b_iterator = db.prefix_iterator(b"bbb"); - assert_eq!(b_iterator.collect::>(), expected) - } + assert_iter(db.prefix_iterator(b"aaa"), &[pair(A1, A1), pair(A2, A2)]); + assert_iter(db.prefix_iterator(b"bbb"), &[pair(B1, B1), pair(B2, B2)]); } } diff --git a/tests/util/mod.rs b/tests/util/mod.rs index b5ff27b..c25118c 100644 --- a/tests/util/mod.rs +++ b/tests/util/mod.rs @@ -1,10 +1,11 @@ +#![allow(dead_code)] + use std::path::{Path, PathBuf}; use rocksdb::{Options, DB}; /// Temporary database path which calls DB::Destroy when DBPath is dropped. pub struct DBPath { - #[allow(dead_code)] dir: tempfile::TempDir, // kept for cleaning up during drop path: PathBuf, } @@ -38,3 +39,27 @@ impl AsRef for &DBPath { &self.path } } + +type Pair = (Box<[u8]>, Box<[u8]>); + +pub fn pair(left: &[u8], right: &[u8]) -> Pair { + (Box::from(left), Box::from(right)) +} + +#[track_caller] +pub fn assert_iter( + iter: rocksdb::DBIteratorWithThreadMode<'_, D>, + want: &[Pair], +) { + assert_eq!(iter.collect::>().as_slice(), want); +} + +#[track_caller] +pub fn assert_iter_reversed( + iter: rocksdb::DBIteratorWithThreadMode<'_, D>, + want: &[Pair], +) { + let mut got = iter.collect::>(); + got.reverse(); + assert_eq!(got.as_slice(), want); +}