Blocktree/crates/btlib-tests/src/fs_queries.rs

use btlib::Result;
use std::fs::{metadata, read_dir, Metadata};
use std::path::Path;

/// Calls the given closure on every entry recursively contained in the given path.
pub fn visit<P: AsRef<Path>, R, F: FnMut(Metadata) -> R>(path: P, visitor: &mut F) -> Result<()> {
    let meta = metadata(&path)?;
    if !meta.is_dir() {
        visitor(meta);
        return Ok(());
    }
    let contents = read_dir(&path)?;
    let mut entry_path = path.as_ref().to_owned();
    entry_path.push("x");
    for entry in contents {
        let entry = entry?;
        entry_path.pop();
        entry_path.push(entry.file_name());
        visit(&entry_path, visitor)?;
    }
    Ok(())
}

/// Recursively sums up the length of each file stored under the given path and returns the total.
/// Note that if the given path contains hard links, the disk usage for the hard linked file will
/// be counted multiple times, once for each link.
pub fn disk_usage<P: AsRef<Path>>(path: P) -> Result<u64> {
    let mut total = 0;
    visit(path, &mut |meta| {
        if meta.is_file() {
            total += meta.len();
        }
    })?;
    Ok(total)
}

pub fn num_files<P: AsRef<Path>>(path: P) -> Result<u64> {
    let mut total = 0;
    visit(path, &mut |meta| {
        if meta.is_file() {
            total += 1;
        }
    })?;
    Ok(total)
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::fs::{create_dir, create_dir_all, write};
    use tempdir::TempDir;

    #[test]
    fn disk_usage_direct_descendants() {
        const BUF: [u8; 8] = [1u8; 8];
        const EXPECTED: u64 = 3 * BUF.len() as u64;
        let dir = TempDir::new("disk_usage").unwrap();
        let dir_path = dir.path();

        for k in 0..3 {
            write(dir_path.join(k.to_string()), &BUF).unwrap();
        }
        let actual = disk_usage(dir_path).unwrap();

        assert_eq!(EXPECTED, actual);
    }

    #[test]
    fn disk_usage_first_gen_dirs() {
        const BUF: [u8; 8] = [1u8; 8];
        const EXPECTED: u64 = 3 * BUF.len() as u64;
        let dir = TempDir::new("disk_usage").unwrap();
        let dir_path = dir.path();

        write(dir_path.join("1"), &BUF).unwrap();
        let sub1 = dir_path.join("sub1");
        create_dir(&sub1).unwrap();
        write(sub1.join("2"), &BUF).unwrap();
        let sub2 = dir_path.join("sub2");
        create_dir(&sub2).unwrap();
        write(sub2.join("3"), &BUF).unwrap();
        let actual = disk_usage(dir_path).unwrap();

        assert_eq!(EXPECTED, actual);
    }

    #[test]
    fn disk_usage_second_gen_dir() {
        const BUF: [u8; 8] = [1u8; 8];
        const EXPECTED: u64 = 3 * BUF.len() as u64;
        let dir = TempDir::new("disk_usage").unwrap();
        let dir_path = dir.path();

        for k in 0..2 {
            write(dir_path.join(k.to_string()), &BUF).unwrap();
        }
        let mut sub = dir_path.to_owned();
        sub.push("sub");
        sub.push("sub");
        create_dir_all(&sub).unwrap();
        sub.push("2");
        write(&sub, &BUF).unwrap();
        let actual = disk_usage(dir_path).unwrap();

        assert_eq!(EXPECTED, actual);
    }

    #[test]
    fn disk_usage_empty_dir() {
        const EXPECTED: u64 = 0;
        let dir = TempDir::new("disk_usage").unwrap();
        let dir_path = dir.path();

        let actual = disk_usage(dir_path).unwrap();

        assert_eq!(EXPECTED, actual);
    }

    #[test]
    fn num_files_direct_descendants() {
        const BUF: [u8; 8] = [1u8; 8];
        const EXPECTED: u64 = 3;
        let dir = TempDir::new("num_files").unwrap();
        let dir_path = dir.path();

        for k in 0..3 {
            write(dir_path.join(k.to_string()), &BUF).unwrap();
        }
        let actual = num_files(dir_path).unwrap();

        assert_eq!(EXPECTED, actual);
    }

    #[test]
    fn num_files_first_gen_dirs() {
        const BUF: [u8; 8] = [1u8; 8];
        const EXPECTED: u64 = 3;
        let dir = TempDir::new("num_files").unwrap();
        let dir_path = dir.path();

        write(dir_path.join("1"), &BUF).unwrap();
        let sub1 = dir_path.join("sub1");
        create_dir(&sub1).unwrap();
        write(sub1.join("2"), &BUF).unwrap();
        let sub2 = dir_path.join("sub2");
        create_dir(&sub2).unwrap();
        write(sub2.join("3"), &BUF).unwrap();
        let actual = num_files(dir_path).unwrap();

        assert_eq!(EXPECTED, actual);
    }

    #[test]
    fn num_files_second_gen_dir() {
        const BUF: [u8; 8] = [1u8; 8];
        const EXPECTED: u64 = 3;
        let dir = TempDir::new("num_files").unwrap();
        let dir_path = dir.path();

        for k in 0..2 {
            write(dir_path.join(k.to_string()), &BUF).unwrap();
        }
        let mut sub = dir_path.to_owned();
        sub.push("sub");
        sub.push("sub");
        create_dir_all(&sub).unwrap();
        sub.push("2");
        write(&sub, &BUF).unwrap();
        let actual = num_files(dir_path).unwrap();

        assert_eq!(EXPECTED, actual);
    }

    #[test]
    fn num_files_empty_dir() {
        const EXPECTED: u64 = 0;
        let dir = TempDir::new("num_files").unwrap();
        let dir_path = dir.path();

        let actual = num_files(dir_path).unwrap();

        assert_eq!(EXPECTED, actual);
    }
}