mithril_aggregator/tools/
file_size.rs

1use anyhow::Context;
2use std::{
3    collections::HashSet,
4    path::{Path, PathBuf},
5};
6
7use mithril_common::StdResult;
8
9/// Compute the size of the given paths that could be files or folders.
10pub(crate) fn compute_size(paths: Vec<PathBuf>) -> StdResult<u64> {
11    fn remove_duplicated_paths(paths: Vec<PathBuf>) -> Vec<PathBuf> {
12        let mut result_folders = vec![];
13        let mut result_files = HashSet::new();
14        let mut paths = paths;
15        paths.sort();
16        for path in paths {
17            if result_folders
18                .last()
19                .is_none_or(|last_folder| !path.starts_with(last_folder))
20            {
21                if path.is_file() {
22                    result_files.insert(path);
23                } else {
24                    result_folders.push(path);
25                }
26            }
27        }
28
29        result_folders.extend(result_files);
30        result_folders
31    }
32
33    let paths = remove_duplicated_paths(paths);
34
35    let mut total = 0;
36    for path_to_include in paths {
37        total += compute_size_of_path(&path_to_include)?;
38    }
39    Ok(total)
40}
41
42/// Compute the size of one given path that could be a file or a folder.
43///
44/// Returns 0 if the path is not a file or a folder.
45pub(crate) fn compute_size_of_path(path: &Path) -> StdResult<u64> {
46    if path.is_file() {
47        let metadata = std::fs::metadata(path)
48            .with_context(|| format!("Failed to read metadata for file: {:?}", path))?;
49
50        return Ok(metadata.len());
51    }
52
53    if path.is_dir() {
54        let entries = std::fs::read_dir(path)
55            .with_context(|| format!("Failed to read directory: {:?}", path))?;
56        let mut directory_size = 0;
57        for entry in entries {
58            let path = entry
59                .with_context(|| format!("Failed to read directory entry in {:?}", path))?
60                .path();
61            directory_size += compute_size_of_path(&path)?;
62        }
63
64        return Ok(directory_size);
65    }
66
67    Ok(0)
68}
69
70#[cfg(test)]
71mod tests {
72    use std::fs::File;
73    use std::io::Write;
74
75    use mithril_common::current_function;
76    use mithril_common::test_utils::TempDir;
77
78    use super::*;
79
80    /// Create a file with the given name in the given dir, write some text to it, and then
81    /// return its path.
82    fn write_dummy_file(optional_size: Option<u64>, dir: &Path, filename: &str) -> PathBuf {
83        let file = dir.join(Path::new(filename));
84        let mut source_file = File::create(&file).unwrap();
85
86        write!(source_file, "This is a test file named '{filename}'").unwrap();
87
88        if let Some(file_size) = optional_size {
89            writeln!(source_file).unwrap();
90            source_file.set_len(file_size).unwrap();
91        }
92
93        file
94    }
95
96    #[test]
97    fn test_compute_file_size() {
98        let test_dir = TempDir::create("utils", current_function!());
99        let file_path = write_dummy_file(Some(4), &test_dir, "file");
100
101        let size = compute_size(vec![file_path]).unwrap();
102        assert_eq!(size, 4);
103    }
104
105    #[test]
106    fn test_compute_multiple_files_size() {
107        let test_dir = TempDir::create("utils", current_function!());
108        let file_path_1 = write_dummy_file(Some(4), &test_dir, "file_1");
109        let file_path_2 = write_dummy_file(Some(7), &test_dir, "file_2");
110
111        let size = compute_size(vec![file_path_1, file_path_2]).unwrap();
112        assert_eq!(size, 11);
113    }
114
115    #[test]
116    fn test_compute_folder_size() {
117        let test_dir = TempDir::create("utils", current_function!());
118        write_dummy_file(Some(4), &test_dir, "file_1");
119        write_dummy_file(Some(7), &test_dir, "file_2");
120
121        let size = compute_size(vec![test_dir]).unwrap();
122        assert_eq!(size, 11);
123    }
124
125    #[test]
126    fn test_compute_multi_folders_size() {
127        let test_dir = TempDir::create("utils", current_function!());
128
129        let sub_dir_1 = test_dir.join("sub_dir_1");
130        std::fs::create_dir(&sub_dir_1).unwrap();
131        write_dummy_file(Some(4), &sub_dir_1, "file_1");
132
133        let sub_dir_2 = test_dir.join("sub_dir_2");
134        std::fs::create_dir(&sub_dir_2).unwrap();
135        write_dummy_file(Some(7), &sub_dir_2, "file_2");
136
137        let sub_dir_3 = test_dir.join("sub_dir_3");
138        std::fs::create_dir(&sub_dir_3).unwrap();
139        write_dummy_file(Some(3), &sub_dir_3, "file_3");
140
141        let size = compute_size(vec![sub_dir_1, sub_dir_2]).unwrap();
142        assert_eq!(size, 11);
143    }
144
145    #[test]
146    fn test_compute_sub_folders_size() {
147        let test_dir = TempDir::create("utils", current_function!());
148
149        let sub_dir_1 = test_dir.join("sub_dir_1");
150        std::fs::create_dir(&sub_dir_1).unwrap();
151        write_dummy_file(Some(4), &sub_dir_1, "file_1");
152
153        let sub_dir_2 = sub_dir_1.join("sub_dir_2");
154        std::fs::create_dir(&sub_dir_2).unwrap();
155        write_dummy_file(Some(7), &sub_dir_2, "file_2");
156
157        let size = compute_size(vec![sub_dir_1]).unwrap();
158        assert_eq!(size, 11);
159    }
160
161    #[test]
162    fn test_compute_size_count_a_file_only_once() {
163        let test_dir = TempDir::create("utils", current_function!());
164        let file_path_1 = write_dummy_file(Some(4), &test_dir, "file_1");
165
166        let size =
167            compute_size(vec![file_path_1.clone(), file_path_1.clone(), file_path_1]).unwrap();
168        assert_eq!(size, 4);
169    }
170
171    #[test]
172    fn test_compute_size_count_a_file_only_once_when_it_s_part_of_a_computed_folder() {
173        let test_dir = TempDir::create("utils", current_function!());
174        let file_path_1 = write_dummy_file(Some(4), &test_dir, "file_1");
175
176        let size = compute_size(vec![test_dir.clone(), file_path_1.clone()]).unwrap();
177        assert_eq!(size, 4);
178
179        let size = compute_size(vec![file_path_1, test_dir]).unwrap();
180        assert_eq!(size, 4);
181    }
182
183    #[test]
184    fn test_compute_file_with_file_name_starting_with_a_folder_name() {
185        let test_dir = TempDir::create("utils", current_function!());
186
187        let sub_dir_1 = test_dir.join("sub_dir_1");
188        std::fs::create_dir(&sub_dir_1).unwrap();
189        write_dummy_file(Some(3), &sub_dir_1, "file_1");
190
191        let file_path_2 = write_dummy_file(Some(4), &test_dir, "sub_dir_1_file_2");
192        assert!(file_path_2
193            .to_str()
194            .unwrap()
195            .starts_with(sub_dir_1.to_str().unwrap()));
196
197        let size = compute_size(vec![sub_dir_1.clone(), file_path_2.clone()]).unwrap();
198        assert_eq!(size, 7);
199
200        let size = compute_size(vec![file_path_2, sub_dir_1]).unwrap();
201        assert_eq!(size, 7);
202    }
203
204    #[test]
205    fn test_compute_file_with_folder_name_starting_with_a_file_name() {
206        let test_dir = TempDir::create("utils", current_function!());
207
208        let file_path_2 = write_dummy_file(Some(4), &test_dir, "file_2");
209
210        let sub_dir_1 = test_dir.join("file_2_sub_dir_1");
211        std::fs::create_dir(&sub_dir_1).unwrap();
212        write_dummy_file(Some(3), &sub_dir_1, "file_1");
213
214        assert!(sub_dir_1
215            .to_str()
216            .unwrap()
217            .starts_with(file_path_2.to_str().unwrap()));
218
219        let size = compute_size(vec![sub_dir_1.clone(), file_path_2.clone()]).unwrap();
220        assert_eq!(size, 7);
221
222        let size = compute_size(vec![file_path_2, sub_dir_1]).unwrap();
223        assert_eq!(size, 7);
224    }
225}