prettyping_rs/engine/
mock.rs

1use std::collections::BTreeMap;
2use std::time::Duration;
3
4use super::{EngineTime, PingEngine, PingEngineError, ProbeRequest, SequenceNumber, TimedEvent};
5
6#[derive(Debug, Clone, Default)]
7pub struct MockEngine {
8    now: EngineTime,
9    sent: Vec<ProbeRequest>,
10    events: BTreeMap<EngineTime, Vec<TimedEvent>>,
11    send_failures: BTreeMap<SequenceNumber, String>,
12}
13
14impl MockEngine {
15    #[must_use]
16    pub fn new() -> Self {
17        Self::default()
18    }
19
20    #[must_use]
21    pub fn with_now(now: Duration) -> Self {
22        Self {
23            now,
24            ..Self::default()
25        }
26    }
27
28    pub fn queue_event(&mut self, event: TimedEvent) {
29        self.events.entry(event.at).or_default().push(event);
30    }
31
32    pub fn queue_events<I>(&mut self, events: I)
33    where
34        I: IntoIterator<Item = TimedEvent>,
35    {
36        for event in events {
37            self.queue_event(event);
38        }
39    }
40
41    pub fn fail_send_for_seq(&mut self, seq: SequenceNumber, message: impl Into<String>) {
42        self.send_failures.insert(seq, message.into());
43    }
44
45    #[must_use]
46    pub fn sent_requests(&self) -> &[ProbeRequest] {
47        &self.sent
48    }
49
50    #[must_use]
51    pub fn now_time(&self) -> EngineTime {
52        self.now
53    }
54}
55
56impl PingEngine for MockEngine {
57    fn now(&self) -> EngineTime {
58        self.now
59    }
60
61    fn send_probe(&mut self, request: ProbeRequest) -> Result<(), PingEngineError> {
62        if request.sent_at != self.now {
63            return Err(PingEngineError::InvalidProbeRequest {
64                seq: request.seq,
65                message: format!(
66                    "sent_at {:?} does not match engine now {:?}",
67                    request.sent_at, self.now
68                ),
69            });
70        }
71
72        if let Some(message) = self.send_failures.remove(&request.seq) {
73            return Err(PingEngineError::SendFailed {
74                seq: request.seq,
75                message,
76            });
77        }
78
79        self.sent.push(request);
80        Ok(())
81    }
82
83    fn poll_until(&mut self, deadline: EngineTime) -> Result<Vec<TimedEvent>, PingEngineError> {
84        if deadline < self.now {
85            return Err(PingEngineError::NonMonotonicPoll);
86        }
87
88        let selected_time = self
89            .events
90            .range(self.now..=deadline)
91            .next()
92            .map(|(event_time, _)| *event_time);
93
94        match selected_time {
95            Some(event_time) => {
96                self.now = event_time;
97                Ok(self.events.remove(&event_time).unwrap_or_default())
98            }
99            None => {
100                self.now = deadline;
101                Ok(Vec::new())
102            }
103        }
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use std::time::Duration;
110
111    use super::MockEngine;
112    use crate::engine::{PingEngine, PingEvent, PingReply, TimedEvent};
113
114    fn ms(value: u64) -> Duration {
115        Duration::from_millis(value)
116    }
117
118    #[test]
119    fn poll_until_ignores_stale_events_and_keeps_time_monotonic() {
120        let mut engine = MockEngine::with_now(ms(1_000));
121        engine.queue_event(TimedEvent {
122            at: ms(900),
123            event: PingEvent::Reply(PingReply::for_seq(1)),
124        });
125
126        let events = engine
127            .poll_until(ms(1_500))
128            .expect("poll should succeed without rewinding");
129
130        assert!(events.is_empty(), "stale event must not be consumed");
131        assert_eq!(engine.now_time(), ms(1_500), "now must remain monotonic");
132    }
133}
prettyping_rs/engine/mock.rs

prettyping_rs/engine/
mock.rs
