1use std::collections::BTreeSet;
5use std::io::Write;
6use std::path::PathBuf;
7
8use crate::Error;
9
10#[derive(Default)]
49pub struct Stager {
50 files: Vec<(&'static str, &'static [u8])>,
51}
52
53impl Stager {
54 #[must_use]
56 pub fn new() -> Self {
57 Self::default()
58 }
59
60 #[must_use]
62 pub fn add(mut self, relative_path: &'static str, bytes: &'static [u8]) -> Self {
63 self.files.push((relative_path, bytes));
64 self
65 }
66
67 #[must_use]
70 pub fn with<I>(mut self, iter: I) -> Self
71 where
72 I: IntoIterator<Item = (&'static str, &'static [u8])>,
73 {
74 self.files.extend(iter);
75 self
76 }
77
78 pub fn stage(self) -> Result<tempfile::TempDir, Error> {
88 let mut seen: BTreeSet<&str> = BTreeSet::new();
89 for (path, _) in &self.files {
90 if !seen.insert(*path) {
91 return Err(Error::DuplicatePath {
92 path: (*path).to_string(),
93 });
94 }
95 }
96
97 let dir = tempfile::tempdir()?;
98 for (rel, bytes) in self.files {
99 let target: PathBuf = dir.path().join(rel);
100 if let Some(parent) = target.parent() {
101 std::fs::create_dir_all(parent)?;
102 }
103 let mut f = std::fs::File::create(&target)?;
104 f.write_all(bytes)?;
105 }
106 Ok(dir)
107 }
108}
109
110#[cfg(test)]
111mod tests {
112 use super::*;
113
114 #[test]
115 fn stages_single_file_at_relative_path() {
116 let dir = Stager::new()
117 .add("a/v1/x.proto", b"syntax = \"proto3\";")
118 .stage()
119 .expect("stage");
120 let p = dir.path().join("a/v1/x.proto");
121 assert!(p.exists(), "expected staged file at {p:?}");
122 let body = std::fs::read(&p).expect("read");
123 assert_eq!(body.as_slice(), b"syntax = \"proto3\";");
124 }
125
126 #[test]
127 fn stages_multiple_files_under_nested_subdirs() {
128 let dir = Stager::new()
129 .add("a/v1/x.proto", b"x")
130 .add("b/v2/y.proto", b"y")
131 .stage()
132 .expect("stage");
133 assert!(dir.path().join("a/v1/x.proto").exists());
134 assert!(dir.path().join("b/v2/y.proto").exists());
135 }
136
137 #[test]
138 fn with_iterator_extends_files() {
139 let pairs: Vec<(&str, &[u8])> = vec![
140 ("a/v1/x.proto", b"x" as &[u8]),
141 ("b/v1/y.proto", b"y" as &[u8]),
142 ];
143 let dir = Stager::new().with(pairs).stage().expect("stage");
144 assert!(dir.path().join("a/v1/x.proto").exists());
145 assert!(dir.path().join("b/v1/y.proto").exists());
146 }
147
148 #[test]
149 fn duplicate_paths_return_error() {
150 let err = Stager::new()
151 .add("a/v1/x.proto", b"first")
152 .add("a/v1/x.proto", b"second")
153 .stage()
154 .expect_err("should error on duplicate");
155 match err {
156 Error::DuplicatePath { path } => assert_eq!(path, "a/v1/x.proto"),
157 other => panic!("expected DuplicatePath, got {other:?}"),
158 }
159 }
160}