Skip to main content

cortexkit_test_support/
daemon.rs

1use crate::{fence::fenced_env_with_roots, fenced_command};
2use std::{
3    ffi::OsStr,
4    path::{Path, PathBuf},
5    process::{ExitStatus, Stdio},
6    time::{Duration, Instant},
7};
8
9/// Command builder whose only spawn operation returns an owned `TestDaemon`.
10/// It enforces scratch-root isolation and keeps the child tied to a lifetime pipe.
11/// Unix consumers require `python3` on PATH to run the process-lifetime watcher.
12pub struct TestDaemonCommand {
13    command: tokio::process::Command,
14    scratch: PathBuf,
15    program: PathBuf,
16    extra_roots: Vec<(String, String)>,
17}
18
19impl TestDaemonCommand {
20    pub fn new(program: &Path, scratch: &Path) -> Self {
21        #[cfg(unix)]
22        let mut command = {
23            let watcher =
24                super::stable_executable("ckdev-daemon-watch", include_bytes!("daemon_watch.py"));
25            let mut command = fenced_command("python3", scratch);
26            command.arg(watcher).arg(program);
27            use std::os::unix::process::CommandExt;
28            command.process_group(0);
29            tokio::process::Command::from(command)
30        };
31        #[cfg(not(unix))]
32        let mut command = tokio::process::Command::from(fenced_command(program, scratch));
33        command.stdin(Stdio::piped()).kill_on_drop(true);
34        Self {
35            command,
36            scratch: scratch.to_owned(),
37            program: program.to_owned(),
38            extra_roots: Vec::new(),
39        }
40    }
41
42    pub fn env(&mut self, key: impl AsRef<OsStr>, value: impl AsRef<OsStr>) -> &mut Self {
43        self.command.env(key, value);
44        self
45    }
46    /// Redirect an additional directory variable into a plain subdirectory of
47    /// scratch. It is reapplied at spawn, so `.env` and `.env_remove` cannot
48    /// bypass it. Panics for a variable the fence already sets (HOME, the XDG
49    /// directories, TMPDIR) and for a directory name that is not a single plain
50    /// path component.
51    pub fn fenced_root(&mut self, variable: &str, directory: &str) -> &mut Self {
52        fenced_env_with_roots(&self.scratch, &[(variable, directory)]);
53        if let Some((_, current)) = self
54            .extra_roots
55            .iter_mut()
56            .find(|(name, _)| name.eq_ignore_ascii_case(variable))
57        {
58            *current = directory.to_owned();
59        } else {
60            self.extra_roots
61                .push((variable.to_owned(), directory.to_owned()));
62        }
63        self
64    }
65    pub fn arg(&mut self, arg: impl AsRef<OsStr>) -> &mut Self {
66        self.command.arg(arg);
67        self
68    }
69    pub fn args<I, S>(&mut self, args: I) -> &mut Self
70    where
71        I: IntoIterator<Item = S>,
72        S: AsRef<OsStr>,
73    {
74        self.command.args(args);
75        self
76    }
77    pub fn env_remove(&mut self, key: impl AsRef<OsStr>) -> &mut Self {
78        self.command.env_remove(key);
79        self
80    }
81    pub fn stdout(&mut self, stdio: impl Into<Stdio>) -> &mut Self {
82        self.command.stdout(stdio);
83        self
84    }
85    pub fn stderr(&mut self, stdio: impl Into<Stdio>) -> &mut Self {
86        self.command.stderr(stdio);
87        self
88    }
89    pub fn kill_on_drop(&mut self, enabled: bool) -> &mut Self {
90        assert!(enabled, "test daemons must die with their harness");
91        self
92    }
93    pub fn spawn(&mut self) -> std::io::Result<TestDaemon> {
94        if !self
95            .program
96            .file_name()
97            .and_then(OsStr::to_str)
98            .is_some_and(|name| name.starts_with("ckdev-"))
99        {
100            return Err(std::io::Error::new(
101                std::io::ErrorKind::InvalidInput,
102                "test daemon program must be named ckdev-*",
103            ));
104        }
105        // Set HOME, the XDG directories, TMPDIR and caller-selected roots last,
106        // just before spawn. These are the variables the daemon reads to decide
107        // where it writes, so a later value always wins over anything a caller
108        // set through `.env`, and the daemon writes only under its scratch tree.
109        let roots: Vec<_> = self
110            .extra_roots
111            .iter()
112            .map(|(name, directory)| (name.as_str(), directory.as_str()))
113            .collect();
114        self.command
115            .envs(fenced_env_with_roots(&self.scratch, &roots));
116        let result = self.command.spawn();
117        #[cfg(unix)]
118        let result = result.map_err(|error| {
119            std::io::Error::new(
120                error.kind(),
121                format!("spawn test daemon watcher python3 (requires python3 on PATH): {error}"),
122            )
123        });
124        let child = result?;
125        Ok(TestDaemon { child })
126    }
127}
128
129/// Owns the parent pipe and process group. Drop is bounded and synchronous, so
130/// when a test panics, the daemon is torn down before the test's scratch
131/// directory is deleted.
132pub struct TestDaemon {
133    child: tokio::process::Child,
134}
135
136impl TestDaemon {
137    pub fn id(&self) -> Option<u32> {
138        self.child.id()
139    }
140    pub fn try_wait(&mut self) -> std::io::Result<Option<ExitStatus>> {
141        self.child.try_wait()
142    }
143    pub fn start_kill(&mut self) -> std::io::Result<()> {
144        self.child.stdin.take();
145        #[cfg(unix)]
146        if let Some(pid) = self.child.id() {
147            // On Unix the spawned child is the watcher script, which leads its own
148            // process group, so this TERM reaches it. It reacts by killing its
149            // daemon and every descendant, including supervised children that
150            // ignore TERM themselves.
151            let group = rustix::process::Pid::from_raw(pid as i32).expect("child PID");
152            let _ = rustix::process::kill_process(group, rustix::process::Signal::TERM);
153        }
154        #[cfg(not(unix))]
155        self.child.start_kill()?;
156        Ok(())
157    }
158    pub async fn wait(&mut self) -> std::io::Result<ExitStatus> {
159        // Tokio closes a child's stdin when waiting. The watcher's stdin is a
160        // lifetime pipe, not input, so keep it owned until the watcher exits.
161        let pipe = self.child.stdin.take();
162        let result = self.child.wait().await;
163        drop(pipe);
164        result
165    }
166    pub async fn kill(&mut self) -> std::io::Result<()> {
167        self.start_kill()?;
168        self.wait().await.map(|_| ())
169    }
170    pub fn stop(&mut self) {
171        let _ = self.start_kill();
172        let deadline = Instant::now() + Duration::from_secs(5);
173        while matches!(self.child.try_wait(), Ok(None)) && Instant::now() < deadline {
174            std::thread::park_timeout(Duration::from_millis(10));
175        }
176        if self.child.id().is_some() {
177            #[cfg(unix)]
178            if let Some(group) = rustix::process::Pid::from_raw(self.child.id().unwrap() as i32) {
179                let _ = rustix::process::kill_process_group(group, rustix::process::Signal::KILL);
180            }
181            let _ = self.child.start_kill();
182        }
183    }
184}
185
186impl Drop for TestDaemon {
187    fn drop(&mut self) {
188        self.stop();
189    }
190}
191
192#[cfg(all(test, unix))]
193mod tests {
194    use super::*;
195
196    #[tokio::test]
197    async fn missing_python3_error_names_the_watcher_requirement() {
198        let scratch = crate::ScratchDir::new("missing-python");
199        let empty_path = scratch.join("empty-path");
200        std::fs::create_dir(&empty_path).unwrap();
201        let program = crate::stable_executable("ckdev-python-requirement", b"#!/bin/sh\nexit 0\n");
202        let result = TestDaemonCommand::new(&program, &scratch)
203            .env("PATH", &empty_path)
204            .spawn();
205        let error = match result {
206            Err(error) => error,
207            Ok(_) => panic!("watcher spawned without python3 on PATH"),
208        };
209        assert_eq!(error.kind(), std::io::ErrorKind::NotFound);
210        let text = error.to_string();
211        assert!(text.contains("python3") && text.contains("PATH"), "{text}");
212    }
213
214    #[tokio::test]
215    async fn spawned_child_sees_enforced_caller_selected_roots() {
216        let scratch = crate::ScratchDir::new("daemon-extra-roots");
217        let output = scratch.join("extra-env");
218        let program = crate::stable_executable(
219            "ckdev-extra-root-probe",
220            b"#!/bin/sh\nprintf '%s\\n' \"$EXAMPLE_STORAGE_ROOT\" \"$EXAMPLE_WORKTREE_ROOT\" > \"$1\"\n",
221        );
222        let mut command = TestDaemonCommand::new(&program, &scratch);
223        command
224            .fenced_root("EXAMPLE_STORAGE_ROOT", "storage")
225            .fenced_root("EXAMPLE_WORKTREE_ROOT", "worktrees")
226            .arg(&output)
227            .env("EXAMPLE_STORAGE_ROOT", "/outside/storage")
228            .env_remove("EXAMPLE_WORKTREE_ROOT");
229        let mut daemon = command.spawn().unwrap();
230        assert!(daemon.wait().await.unwrap().success());
231        let expected = ["storage", "worktrees"]
232            .map(|dir| scratch.join(dir).display().to_string())
233            .join("\n")
234            + "\n";
235        assert_eq!(std::fs::read_to_string(output).unwrap(), expected);
236        assert!(scratch.join("storage").is_dir() && scratch.join("worktrees").is_dir());
237    }
238
239    #[tokio::test]
240    async fn non_ckdev_program_is_refused() {
241        let scratch = crate::ScratchDir::new("daemon-refusal");
242        let program = crate::stable_executable("not-a-dev-daemon", b"#!/bin/sh\nexit 0\n");
243        let result = TestDaemonCommand::new(&program, &scratch).spawn();
244        assert!(
245            matches!(result, Err(ref error) if error.kind()==std::io::ErrorKind::InvalidInput),
246            "non-ckdev daemon was not refused"
247        );
248    }
249
250    #[tokio::test]
251    async fn spawned_child_sees_enforced_xdg_roots() {
252        let scratch = crate::ScratchDir::new("daemon-xdg");
253        let output = scratch.join("xdg-env");
254        let program = crate::stable_executable("ckdev-xdg-probe", b"#!/bin/sh\nprintf '%s\\n' \"$XDG_DATA_HOME\" \"$XDG_RUNTIME_DIR\" \"$XDG_CONFIG_HOME\" > \"$1\"\n");
255        let mut command = TestDaemonCommand::new(&program, &scratch);
256        command
257            .arg(&output)
258            .env("XDG_DATA_HOME", "/outside/data")
259            .env_remove("XDG_RUNTIME_DIR")
260            .env("XDG_CONFIG_HOME", "/outside/config");
261        let mut daemon = command.spawn().unwrap();
262        assert!(daemon.wait().await.unwrap().success());
263        let expected = ["data", "runtime", "config"]
264            .map(|d| scratch.join(d).display().to_string())
265            .join("\n")
266            + "\n";
267        assert_eq!(std::fs::read_to_string(output).unwrap(), expected);
268    }
269
270    #[tokio::test]
271    #[should_panic(expected = "test daemons must die with their harness")]
272    async fn kill_on_drop_false_is_refused() {
273        let scratch = crate::ScratchDir::new("daemon-kill-guard");
274        TestDaemonCommand::new(Path::new("ckdev-fixture"), &scratch).kill_on_drop(false);
275    }
276
277    fn alive(pid: i32) -> bool {
278        subc_os::process_identity::liveness(pid as u32, None)
279            == subc_os::process_identity::Liveness::Alive
280    }
281
282    async fn fixture() -> (super::super::ScratchDir, TestDaemon, i32) {
283        let scratch = super::super::ScratchDir::new("daemon-lifetime");
284        let pid_file = scratch.join("daemon.pid");
285        let program = super::super::stable_executable(
286            "ckdev-lifetime-fixture",
287            b"#!/bin/sh\ntrap '' TERM\nexec sleep 600\n",
288        );
289        let daemon = TestDaemonCommand::new(&program, scratch.path())
290            .env("CORTEXKIT_TEST_DAEMON_PID_FILE", &pid_file)
291            .spawn()
292            .unwrap();
293        let deadline = Instant::now() + Duration::from_secs(5);
294        while !pid_file.exists() {
295            assert!(Instant::now() < deadline, "fixture daemon did not start");
296            tokio::time::sleep(Duration::from_millis(10)).await;
297        }
298        let pid = std::fs::read_to_string(pid_file).unwrap().parse().unwrap();
299        assert!(alive(pid), "positive control: daemon is running");
300        (scratch, daemon, pid)
301    }
302
303    async fn assert_gone(pid: i32) {
304        let deadline = Instant::now() + Duration::from_secs(5);
305        while alive(pid) && Instant::now() < deadline {
306            tokio::time::sleep(Duration::from_millis(10)).await;
307        }
308        assert!(!alive(pid), "dropped harness leaked daemon pid {pid}");
309    }
310
311    #[tokio::test]
312    async fn dropping_harness_mid_run_reaps_its_daemon() {
313        let (_scratch, daemon, pid) = fixture().await;
314        drop(daemon);
315        assert_gone(pid).await;
316    }
317
318    #[tokio::test]
319    async fn panic_unwinding_reaps_its_daemon() {
320        let (_scratch, daemon, pid) = fixture().await;
321        let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
322            let _daemon = daemon;
323            panic!("intentional harness assertion");
324        }));
325        assert!(panic.is_err());
326        assert_gone(pid).await;
327    }
328
329    #[tokio::test]
330    async fn parent_pipe_eof_reaps_its_daemon_without_drop() {
331        let (_scratch, mut daemon, pid) = fixture().await;
332        daemon.child.stdin.take();
333        tokio::time::timeout(Duration::from_secs(5), daemon.wait())
334            .await
335            .unwrap()
336            .unwrap();
337        assert_gone(pid).await;
338    }
339
340    #[test]
341    fn fleet_binary_is_copied_to_ckdev_scratch_before_exec() {
342        let scratch = super::super::ScratchDir::new("fleet-name");
343        let source = scratch.join("ck-name-probe");
344        std::fs::write(&source, b"#!/bin/sh\nprintf '%s\\n' \"$1\"\n").unwrap();
345        use std::os::unix::fs::PermissionsExt;
346        std::fs::set_permissions(&source, std::fs::Permissions::from_mode(0o755)).unwrap();
347        let staged = crate::stage_test_binary(&source);
348        assert_eq!(staged.file_name().unwrap(), "ckdev-name-probe");
349        assert!(staged.starts_with("/tmp"));
350        assert_eq!(
351            std::fs::read(source).unwrap(),
352            std::fs::read(&staged).unwrap()
353        );
354        let output = super::super::fenced_command(&staged, scratch.path())
355            .arg("live")
356            .output()
357            .unwrap();
358        assert!(output.status.success());
359        assert_eq!(output.stdout, b"live\n");
360    }
361}