add some tests

crates/core/src/commands/check.rs

@@ -8,7 +8,7 @@ use bytes::Bytes;
 use bytesize::ByteSize;
 use derive_setters::Setters;
 use log::{debug, error, warn};
-use rand::{prelude::SliceRandom, thread_rng};
+use rand::{prelude::SliceRandom, thread_rng, Rng};
 use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 use zstd::stream::decode_all;
 
@@ -44,9 +44,50 @@ pub enum ReadSubset {
     IdSubSet((u32, u32)),
 }
 
-fn id_matches_n_m(id: &Id, n: u32, m: u32) -> bool {
-    let short_id: u32 = id.transmute();
-    short_id % m == n % m
+impl ReadSubset {
+    fn apply(self, packs: impl IntoIterator<Item = IndexPack>) -> Vec<IndexPack> {
+        self.apply_with_rng(packs, &mut thread_rng())
+    }
+    fn apply_with_rng(
+        self,
+        packs: impl IntoIterator<Item = IndexPack>,
+        rng: &mut impl Rng,
+    ) -> Vec<IndexPack> {
+        fn id_matches_n_m(id: &Id, n: u32, m: u32) -> bool {
+            let short_id: u32 = id.transmute();
+            short_id % m == n % m
+        }
+
+        let mut total_size: u64 = 0;
+        let mut packs: Vec<_> = packs
+            .into_iter()
+            .inspect(|p| total_size += u64::from(p.pack_size()))
+            .collect();
+
+        // Apply read-subset option
+        if let Some(mut size) = match self {
+            ReadSubset::All => None,
+            ReadSubset::Percentage(p) => Some((total_size as f64 * p / 100.0) as u64),
+            ReadSubset::Size(s) => Some(s),
+            ReadSubset::IdSubSet((n, m)) => {
+                packs.retain(|p| id_matches_n_m(&p.id, n, m));
+                None
+            }
+        } {
+            // random subset of given size is required
+            packs.shuffle(rng);
+            packs.retain(|p| {
+                let p_size = u64::from(p.pack_size());
+                if size > p_size {
+                    size = size.saturating_sub(p_size);
+                    true
+                } else {
+                    false
+                }
+            });
+        }
+        packs
+    }
 }
 
 impl FromStr for ReadSubset {
@@ -164,45 +205,18 @@ impl CheckOptions {
         let packs = check_trees(be, &index_be, trees, pb)?;
 
         if self.read_data {
-            let mut total_size: u64 = 0;
-            let mut packs: Vec<_> = index_be
+            let packs = index_be
                 .into_index()
                 .into_iter()
-                .filter(|p| packs.contains(&p.id))
-                .map(|p| {
-                    let size = u64::from(p.pack_size());
-                    (p, size)
-                })
-                .inspect(|(_, size)| total_size += size)
-                .collect();
+                .filter(|p| packs.contains(&p.id));
 
-            // Apply read-subset option
-            if let Some(mut size) = match self.read_data_subset {
-                ReadSubset::All => None,
-                ReadSubset::Percentage(p) => Some((total_size as f64 * p / 100.0) as u64),
-                ReadSubset::Size(s) => Some(s),
-                ReadSubset::IdSubSet((n, m)) => {
-                    packs.retain(|(p, _)| id_matches_n_m(&p.id, n, m));
-                    None
-                }
-            } {
-                // random subset of given size is required
-                packs.shuffle(&mut thread_rng());
-                packs.retain(|(_, p_size)| {
-                    if size > *p_size {
-                        size = size.saturating_sub(*p_size);
-                        true
-                    } else {
-                        false
-                    }
-                });
-            }
+            let packs = self.read_data_subset.apply(packs);
 
-            let total_pack_size = packs.iter().map(|(_, size)| size).sum();
+            let total_pack_size = packs.iter().map(|pack| u64::from(pack.pack_size())).sum();
             let p = pb.progress_bytes("reading pack data...");
             p.set_length(total_pack_size);
 
-            packs.into_par_iter().for_each(|(pack, _)| {
+            packs.into_par_iter().for_each(|pack| {
                 let id = pack.id;
                 let data = be.read_full(FileType::Pack, &id).unwrap();
                 match check_pack(be, pack, data, &p) {
@@ -603,3 +617,77 @@ fn check_pack(
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use insta::assert_ron_snapshot;
+    use rand::{rngs::StdRng, Rng, SeedableRng};
+    use rstest::{fixture, rstest};
+
+    const PACK_SIZE: u32 = 100_000_000;
+
+    #[fixture]
+    fn rng() -> StdRng {
+        StdRng::seed_from_u64(5)
+    }
+    fn test_packs(rng: &mut impl Rng) -> Vec<IndexPack> {
+        (0..500)
+            .map(|_| IndexPack {
+                id: PackId::from(Id::random_from_rng(rng)),
+                blobs: Vec::new(),
+                time: None,
+                size: Some(rng.gen_range(0..PACK_SIZE)),
+            })
+            .collect()
+    }
+
+    #[rstest]
+    #[case("all")]
+    #[case("5/12")]
+    #[case("5%")]
+    #[case("250MiB")]
+    fn test_read_subset(mut rng: StdRng, #[case] s: &str) {
+        let size =
+            |packs: &[IndexPack]| -> u64 { packs.iter().map(|p| u64::from(p.pack_size())).sum() };
+
+        let test_packs = test_packs(&mut rng);
+        let total_size = size(&test_packs);
+
+        let subset: ReadSubset = s.parse().unwrap();
+        let packs = subset.apply_with_rng(test_packs, &mut rng);
+        let test_size = size(&packs);
+
+        match subset {
+            ReadSubset::All => assert_eq!(test_size, total_size),
+            ReadSubset::Percentage(s) => assert!(test_size <= (total_size as f64 * s) as u64),
+            ReadSubset::Size(size) => {
+                assert!(test_size < size && size < test_size + PACK_SIZE as u64)
+            }
+            _ => {}
+        };
+
+        assert_ron_snapshot!(s, packs);
+    }
+
+    fn test_read_subset_n_m() {
+        let test_packs = test_packs(&mut thread_rng());
+        let mut all_packs: BTreeSet<_> = test_packs.iter().map(|pack| pack.id).collect();
+
+        let mut remove = |s: &str| {
+            let subset: ReadSubset = s.parse().unwrap();
+            let packs = subset.apply(test_packs.clone());
+            for pack in packs {
+                assert!(all_packs.remove(&pack.id));
+            }
+        };
+
+        remove("1/5");
+        remove("2/5");
+        remove("3/5");
+        remove("4/5");
+        remove("5/5");
+
+        assert!(all_packs.is_empty());
+    }
+}

crates/core/src/commands/snapshots/rustic_core__commands__check__tests__250MiB.snap

@@ -0,0 +1,66 @@
+---
+source: crates/core/src/commands/check.rs
+expression: packs
+---
+[
+  IndexPack(
+    id: Id("50ec36b55a9b35de779a2757571d2e4b16f6d1ca3ff73ade120901262c2b265d"),
+    blobs: [],
+    size: Some(25543042),
+  ),
+  IndexPack(
+    id: Id("d07f392b1dfebfdf50bb10a4a7857b543cbe246f84549af8902dfde9737fd425"),
+    blobs: [],
+    size: Some(16489215),
+  ),
+  IndexPack(
+    id: Id("f879520cb3e653a0c77dcc37dcd813a71ec7d972a4a6185b13f2ed240c6e72e1"),
+    blobs: [],
+    size: Some(3342498),
+  ),
+  IndexPack(
+    id: Id("8ba2cc8cf44054f35be413f12ba83969deb01f1c44660822cee1b960d69a7526"),
+    blobs: [],
+    size: Some(46887261),
+  ),
+  IndexPack(
+    id: Id("e6bae2c9c6b9d8b8a72b45590ffc6c8e034083d6a8180877f6d270537a1ac214"),
+    blobs: [],
+    size: Some(47685315),
+  ),
+  IndexPack(
+    id: Id("f77e206f69693ae3490de38ce00f5e89ae7db4808b770c60e07d64815ee0478d"),
+    blobs: [],
+    size: Some(9217773),
+  ),
+  IndexPack(
+    id: Id("20b51c8c49aff07d7063c76a863cbdcea845989ef79d4a3f8ff599687eaebe48"),
+    blobs: [],
+    size: Some(83677501),
+  ),
+  IndexPack(
+    id: Id("9630b7b1e6329e7c28eb0eeb4e0df36bbf45acf3ba5de4a0403b77e47216857a"),
+    blobs: [],
+    size: Some(24144078),
+  ),
+  IndexPack(
+    id: Id("33e053041d2de235e03cc219a8b8300d8f1e35ee034c45f4613ea782d5e672f2"),
+    blobs: [],
+    size: Some(2254122),
+  ),
+  IndexPack(
+    id: Id("7583c1099bf604771a03af7627f4122a59da07db7358484b8543e881a7939b3f"),
+    blobs: [],
+    size: Some(68829),
+  ),
+  IndexPack(
+    id: Id("a94f61701a165181c6940584ca0cd2c2355e5e1eb65a3a295fc4d1c02fa81138"),
+    blobs: [],
+    size: Some(2347337),
+  ),
+  IndexPack(
+    id: Id("5ae22a813d32049b56ac2760a5a34b8f66e30b5232f66bb8eae420c7022197e8"),
+    blobs: [],
+    size: Some(452910),
+  ),
+]