Skip to main content

term_test_support/
guard.rs

1//! Generic RAII cleanup guard for test-owned resources.
2
3use std::fmt;
4
5/// Runs a cleanup closure exactly once when dropped, including during panic
6/// unwinding. Use it to guarantee spawned processes are killed, PTYs are
7/// closed, or temp state is removed even when an assertion fails mid-test.
8///
9/// `FnOnce` (not `FnMut`) is sufficient and most permissive: the guard fires
10/// exactly one time, so the closure may move its captured state.
11///
12/// Call [`defuse`](KillOnDrop::defuse) to consume the guard without running
13/// the cleanup (e.g. when the happy path already performed a graceful
14/// teardown).
15pub struct KillOnDrop<F: FnOnce()> {
16    cleanup: Option<F>,
17}
18
19impl<F: FnOnce()> KillOnDrop<F> {
20    /// Arm the guard with `cleanup`.
21    pub fn new(cleanup: F) -> Self {
22        Self {
23            cleanup: Some(cleanup),
24        }
25    }
26
27    /// Consume the guard WITHOUT running the cleanup closure.
28    pub fn defuse(mut self) {
29        self.cleanup = None;
30    }
31}
32
33impl<F: FnOnce()> Drop for KillOnDrop<F> {
34    fn drop(&mut self) {
35        if let Some(cleanup) = self.cleanup.take() {
36            cleanup();
37        }
38    }
39}
40
41impl<F: FnOnce()> fmt::Debug for KillOnDrop<F> {
42    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43        f.debug_struct("KillOnDrop")
44            .field("armed", &self.cleanup.is_some())
45            .finish()
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use std::sync::Arc;
53    use std::sync::atomic::{AtomicU64, Ordering};
54
55    #[test]
56    fn cleanup_runs_on_drop() {
57        let counter = Arc::new(AtomicU64::new(0));
58        let seen = Arc::clone(&counter);
59        {
60            let _guard = KillOnDrop::new(move || {
61                seen.fetch_add(1, Ordering::Release);
62            });
63            assert_eq!(counter.load(Ordering::Acquire), 0);
64        }
65        assert_eq!(counter.load(Ordering::Acquire), 1);
66    }
67
68    #[test]
69    fn defused_guard_skips_cleanup() {
70        let counter = Arc::new(AtomicU64::new(0));
71        let seen = Arc::clone(&counter);
72        {
73            let guard = KillOnDrop::new(move || {
74                seen.fetch_add(1, Ordering::Release);
75            });
76            guard.defuse();
77        }
78        assert_eq!(counter.load(Ordering::Acquire), 0);
79    }
80
81    #[test]
82    fn cleanup_runs_during_panic_unwind() {
83        let counter = Arc::new(AtomicU64::new(0));
84        let seen = Arc::clone(&counter);
85        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
86            let _guard = KillOnDrop::new(|| {
87                seen.fetch_add(1, Ordering::Release);
88            });
89            panic!("simulated assertion failure");
90        }));
91        assert!(result.is_err());
92        assert_eq!(counter.load(Ordering::Acquire), 1);
93    }
94}