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