ic_testkit/pic/
time.rs

1use pocket_ic::PocketIc;
2use std::{fmt, task::Poll, time::Duration};
3
4/// A caller-owned readiness check failed or exhausted its simulated round budget.
5#[derive(Clone, Debug, Eq, PartialEq)]
6pub enum TickUntilError<E> {
7    /// The predicate was still pending after this many advance/tick rounds.
8    ProgressLimit { rounds: u32 },
9    /// The predicate returned its own failure.
10    Failed(E),
11}
12
13impl<E: fmt::Display> fmt::Display for TickUntilError<E> {
14    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
15        match self {
16            Self::ProgressLimit { rounds } => {
17                write!(
18                    formatter,
19                    "readiness remained pending after {rounds} rounds"
20                )
21            }
22            Self::Failed(error) => write!(formatter, "readiness check failed: {error}"),
23        }
24    }
25}
26
27impl<E: std::error::Error + 'static> std::error::Error for TickUntilError<E> {
28    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
29        match self {
30            Self::ProgressLimit { .. } => None,
31            Self::Failed(error) => Some(error),
32        }
33    }
34}
35
36/// Check a caller-owned predicate, advancing simulated time and ticking while pending.
37///
38/// The predicate runs once before any mutation. Each pending result permits one
39/// round of `advance_time(advance)`, then `tick()`, then another predicate check,
40/// up to `max_rounds`. A zero budget still performs the initial check. A zero
41/// duration permits ticks without advancing time. Completion on the last allowed
42/// round succeeds; predicate failures return immediately without another round.
43///
44/// This bounds simulated progression, not wall-clock time spent inside PocketIC
45/// calls or the predicate. PocketIC operation failures retain upstream panic
46/// behavior. Applications own readiness meaning and predicate error types.
47pub fn tick_until<T, E>(
48    pic: &PocketIc,
49    max_rounds: u32,
50    advance: Duration,
51    mut poll: impl FnMut(&PocketIc) -> Poll<Result<T, E>>,
52) -> Result<T, TickUntilError<E>> {
53    poll_with_rounds(
54        max_rounds,
55        || {
56            pic.advance_time(advance);
57            pic.tick();
58        },
59        || poll(pic),
60    )
61}
62
63fn poll_with_rounds<T, E>(
64    max_rounds: u32,
65    mut round: impl FnMut(),
66    mut poll: impl FnMut() -> Poll<Result<T, E>>,
67) -> Result<T, TickUntilError<E>> {
68    let mut rounds = 0;
69    loop {
70        match poll() {
71            Poll::Ready(result) => return result.map_err(TickUntilError::Failed),
72            Poll::Pending if rounds == max_rounds => {
73                return Err(TickUntilError::ProgressLimit { rounds });
74            }
75            Poll::Pending => {
76                round();
77                rounds += 1;
78            }
79        }
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use std::cell::Cell;
87
88    #[test]
89    fn initial_completion_and_failure_do_not_advance() {
90        for budget in [0, 3] {
91            for result in [Ok(42), Err("predicate failure")] {
92                assert_eq!(
93                    poll_with_rounds(
94                        budget,
95                        || panic!("unexpected round"),
96                        || Poll::Ready(result)
97                    ),
98                    result.map_err(TickUntilError::Failed)
99                );
100            }
101        }
102    }
103
104    #[test]
105    fn completion_on_last_round_preserves_poll_order() {
106        let rounds = Cell::new(0);
107        let polls = Cell::new(0);
108        let result = poll_with_rounds(
109            3,
110            || {
111                assert_eq!(polls.get(), rounds.get() + 1);
112                rounds.set(rounds.get() + 1);
113            },
114            || {
115                assert_eq!(polls.get(), rounds.get());
116                polls.set(polls.get() + 1);
117                if rounds.get() == 3 {
118                    Poll::Ready(Ok::<_, ()>(42))
119                } else {
120                    Poll::Pending
121                }
122            },
123        );
124        assert_eq!(result, Ok(42));
125        assert_eq!((rounds.get(), polls.get()), (3, 4));
126    }
127
128    #[test]
129    fn exhaustion_performs_exactly_the_allowed_rounds() {
130        for budget in [0, 1, 3] {
131            let rounds = Cell::new(0);
132            let polls = Cell::new(0);
133            let result = poll_with_rounds(
134                budget,
135                || rounds.set(rounds.get() + 1),
136                || {
137                    polls.set(polls.get() + 1);
138                    Poll::<Result<(), ()>>::Pending
139                },
140            );
141            assert_eq!(
142                result,
143                Err(TickUntilError::ProgressLimit { rounds: budget })
144            );
145            assert_eq!((rounds.get(), polls.get()), (budget, budget + 1));
146        }
147    }
148
149    #[test]
150    fn failure_after_progress_preserves_the_original_cause() {
151        let rounds = Cell::new(0);
152        let result = poll_with_rounds(
153            3,
154            || rounds.set(rounds.get() + 1),
155            || {
156                if rounds.get() == 1 {
157                    Poll::Ready(Err::<(), _>(std::io::Error::from(
158                        std::io::ErrorKind::PermissionDenied,
159                    )))
160                } else {
161                    Poll::Pending
162                }
163            },
164        );
165        let error = result.unwrap_err();
166        let source = std::error::Error::source(&error).unwrap();
167        assert_eq!(
168            source.downcast_ref::<std::io::Error>().unwrap().kind(),
169            std::io::ErrorKind::PermissionDenied
170        );
171        assert_eq!(rounds.get(), 1);
172    }
173}
ic_testkit/pic/time.rs

ic_testkit/pic/
time.rs
