Skip to main content

csusage_test_support/
lib.rs

1use std::{
2    ffi::{OsStr, OsString},
3    fs,
4    path::{Path, PathBuf},
5    sync::{Mutex, MutexGuard},
6};
7
8use assert_fs::{
9    TempDir,
10    fixture::{ChildPath, FileWriteStr, PathChild, PathCreateDir},
11};
12
13pub mod claude_science;
14pub mod dsh;
15pub mod openhands;
16pub mod zcode;
17
18static ENV_LOCK: Mutex<()> = Mutex::new(());
19
20fn env_lock() -> MutexGuard<'static, ()> {
21    ENV_LOCK.lock().unwrap_or_else(|error| error.into_inner())
22}
23
24pub struct EnvVarGuard {
25    key: &'static str,
26    previous: Option<OsString>,
27    _guard: MutexGuard<'static, ()>,
28}
29
30impl EnvVarGuard {
31    pub fn set(key: &'static str, value: impl AsRef<OsStr>) -> Self {
32        let guard = env_lock();
33        let previous = std::env::var_os(key);
34        unsafe { std::env::set_var(key, value) };
35        Self {
36            key,
37            previous,
38            _guard: guard,
39        }
40    }
41}
42
43impl Drop for EnvVarGuard {
44    fn drop(&mut self) {
45        match self.previous.take() {
46            Some(value) => unsafe { std::env::set_var(self.key, value) },
47            None => unsafe { std::env::remove_var(self.key) },
48        }
49    }
50}
51
52pub struct EnvVarsGuard {
53    previous: Vec<(&'static str, Option<OsString>)>,
54    _guard: MutexGuard<'static, ()>,
55}
56
57impl EnvVarsGuard {
58    pub fn set_many(vars: impl IntoIterator<Item = (&'static str, Option<OsString>)>) -> Self {
59        let guard = env_lock();
60        let mut previous = Vec::new();
61        for (key, value) in vars {
62            previous.push((key, std::env::var_os(key)));
63            match value {
64                Some(value) => unsafe { std::env::set_var(key, value) },
65                None => unsafe { std::env::remove_var(key) },
66            }
67        }
68        Self {
69            previous,
70            _guard: guard,
71        }
72    }
73}
74
75impl Drop for EnvVarsGuard {
76    fn drop(&mut self) {
77        for (key, value) in self.previous.drain(..).rev() {
78            match value {
79                Some(value) => unsafe { std::env::set_var(key, value) },
80                None => unsafe { std::env::remove_var(key) },
81            }
82        }
83    }
84}
85
86pub struct Fixture {
87    dir: TempDir,
88}
89
90impl Fixture {
91    pub fn new() -> Self {
92        Self {
93            dir: TempDir::new().expect("failed to create temporary fixture directory"),
94        }
95    }
96
97    pub fn root(&self) -> &Path {
98        self.dir.path()
99    }
100
101    #[must_use]
102    pub fn path(&self, path: impl AsRef<Path>) -> PathBuf {
103        self.dir.path().join(path)
104    }
105
106    fn child(&self, path: impl AsRef<Path>) -> ChildPath {
107        self.dir.child(path)
108    }
109
110    #[must_use]
111    pub fn write_file(&self, path: impl AsRef<Path>, contents: impl AsRef<str>) -> PathBuf {
112        let child = self.child(path);
113        if let Some(parent) = child.path().parent() {
114            fs::create_dir_all(parent).expect("failed to create fixture file parent directory");
115        }
116        child
117            .write_str(contents.as_ref())
118            .expect("failed to write fixture file");
119        child.path().to_path_buf()
120    }
121
122    #[must_use]
123    pub fn create_dir_all(&self, path: impl AsRef<Path>) -> PathBuf {
124        let child = self.child(path);
125        child
126            .create_dir_all()
127            .expect("failed to create fixture directory");
128        child.path().to_path_buf()
129    }
130}
131
132impl Default for Fixture {
133    fn default() -> Self {
134        Self::new()
135    }
136}
137
138#[macro_export]
139macro_rules! fs_fixture {
140    ({ $($path:literal : $contents:expr_2021),* $(,)? }) => {{
141        let fixture = $crate::Fixture::new();
142        $(
143            let _ = fixture.write_file($path, $contents);
144        )*
145        fixture
146    }};
147}
148
149#[cfg(test)]
150mod tests {
151    #[test]
152    fn creates_inline_fixture_tree() {
153        let fixture = fs_fixture!({
154            "projects/example/session.jsonl": "{}\n",
155        });
156
157        assert_eq!(
158            std::fs::read_to_string(fixture.path("projects/example/session.jsonl")).unwrap(),
159            "{}\n"
160        );
161    }
162
163    #[test]
164    fn creates_incremental_fixture_tree() {
165        let fixture = fs_fixture!({});
166        let _ = fixture.write_file("projects/example/session/chat.jsonl", "{}\n");
167
168        assert!(
169            fixture
170                .path("projects/example/session/chat.jsonl")
171                .is_file()
172        );
173    }
174}