polyhorn_cli/test/
output.rs

1use fs3::FileExt;
2use std::collections::HashMap;
3use std::fs::{File, OpenOptions};
4use std::io::Write;
5use std::path::{Path, PathBuf};
6
7use super::Metadata;
8
9/// Represents an output directory on disk.
10pub struct Output {
11    path: PathBuf,
12    index_file: File,
13    all_metadata: HashMap<String, Metadata>,
14}
15
16impl Output {
17    /// Creates a new output directory with the given path.
18    pub fn new<P>(path: P) -> Output
19    where
20        P: AsRef<Path>,
21    {
22        let path = path.as_ref().to_path_buf();
23        let _ = std::fs::create_dir_all(&path);
24        let index_file = OpenOptions::new()
25            .create(true)
26            .write(true)
27            .open(path.join("Snapshots.toml"))
28            .unwrap();
29
30        index_file.lock_exclusive().unwrap();
31
32        Output {
33            path,
34            index_file,
35            all_metadata: HashMap::new(),
36        }
37    }
38
39    /// Stores a snapshot with the given metadata in the output directory.
40    pub fn store(&mut self, metadata: Metadata, screenshot: Vec<u8>) {
41        let digest = metadata.digest().to_string();
42        self.all_metadata.insert(digest.clone(), metadata);
43
44        File::create(self.path.join(digest + ".png"))
45            .unwrap()
46            .write_all(&screenshot)
47            .unwrap();
48    }
49}
50
51impl Drop for Output {
52    fn drop(&mut self) {
53        self.index_file.set_len(0).unwrap();
54        self.index_file
55            .write_all(
56                toml::to_string_pretty(&self.all_metadata)
57                    .unwrap()
58                    .as_bytes(),
59            )
60            .unwrap();
61
62        self.index_file.unlock().unwrap();
63    }
64}