1use std::{cell::RefCell, future::Future, rc::Rc, time::Duration};
8
9use futures::{channel::oneshot, task::SpawnError};
10use lenso_kernel::{DeterministicDriver, DriverTask, RuntimeDriver};
11
12use crate::FaultInjector;
13
14#[derive(Clone, Debug, Default)]
16pub struct TestSimulator {
17 driver: DeterministicDriver,
18 faults: FaultInjector,
19}
20
21impl TestSimulator {
22 pub fn new() -> Self {
24 Self::default()
25 }
26
27 pub fn gate(&self, name: impl Into<String>) -> SimulatorGate {
29 SimulatorGate {
30 name: name.into(),
31 state: Rc::new(RefCell::new(GateState::default())),
32 }
33 }
34
35 pub fn resource(&self, name: impl Into<String>) -> SimulatorResource {
42 SimulatorResource {
43 name: name.into(),
44 state: Rc::new(RefCell::new(ResourceState::default())),
45 }
46 }
47
48 pub fn faults(&self) -> FaultInjector {
54 self.faults.clone()
55 }
56
57 pub fn run<F: Future>(&self, future: F) -> F::Output {
63 self.driver.run(future)
64 }
65
66 pub fn pump(&self) {
72 self.run(self.driver.yield_now());
73 }
74
75 pub fn spawn<F>(&self, future: F) -> Result<DriverTask, SpawnError>
77 where
78 F: Future<Output = ()> + 'static,
79 {
80 self.driver.spawn_local(Box::pin(future))
81 }
82
83 pub fn advance(&self, duration: Duration) {
85 self.driver.advance(duration);
86 }
87
88 pub fn now(&self) -> Duration {
90 self.driver.now()
91 }
92
93 pub(crate) fn driver(&self) -> DeterministicDriver {
94 self.driver.clone()
95 }
96}
97
98#[derive(Debug, Default)]
99struct ResourceState {
100 acquisitions: usize,
101 frozen: bool,
102 thaw_notifiers: Vec<oneshot::Sender<()>>,
103}
104
105#[derive(Clone, Debug)]
111pub struct SimulatorResource {
112 name: String,
113 state: Rc<RefCell<ResourceState>>,
114}
115
116impl SimulatorResource {
117 pub fn name(&self) -> &str {
119 &self.name
120 }
121
122 pub async fn acquire(&self) {
124 loop {
125 let thaw = {
126 let mut state = self.state.borrow_mut();
127 if state.frozen {
128 let (sender, receiver) = oneshot::channel();
129 state.thaw_notifiers.push(sender);
130 Some(receiver)
131 } else {
132 state.acquisitions += 1;
133 None
134 }
135 };
136
137 let Some(thaw) = thaw else {
138 return;
139 };
140 let _ = thaw.await;
141 }
142 }
143
144 #[must_use]
148 pub fn freeze(&self) -> bool {
149 let mut state = self.state.borrow_mut();
150 if state.frozen {
151 return false;
152 }
153 state.frozen = true;
154 true
155 }
156
157 #[must_use]
161 pub fn thaw(&self) -> bool {
162 let notifiers = {
163 let mut state = self.state.borrow_mut();
164 if !state.frozen {
165 return false;
166 }
167 state.frozen = false;
168 std::mem::take(&mut state.thaw_notifiers)
169 };
170 for notifier in notifiers {
171 let _ = notifier.send(());
172 }
173 true
174 }
175
176 pub fn acquisition_count(&self) -> usize {
178 self.state.borrow().acquisitions
179 }
180
181 pub fn is_frozen(&self) -> bool {
183 self.state.borrow().frozen
184 }
185}
186
187#[derive(Debug, Default)]
188struct GateState {
189 reached: usize,
190 released: bool,
191 reached_notifiers: Vec<oneshot::Sender<()>>,
192 release_notifiers: Vec<oneshot::Sender<()>>,
193}
194
195#[derive(Clone, Debug)]
197pub struct SimulatorGate {
198 name: String,
199 state: Rc<RefCell<GateState>>,
200}
201
202impl SimulatorGate {
203 pub fn name(&self) -> &str {
205 &self.name
206 }
207
208 pub async fn wait(&self) {
210 let release = {
211 let mut state = self.state.borrow_mut();
212 state.reached += 1;
213 for notifier in state.reached_notifiers.drain(..) {
214 let _ = notifier.send(());
215 }
216 if state.released {
217 None
218 } else {
219 let (sender, receiver) = oneshot::channel();
220 state.release_notifiers.push(sender);
221 Some(receiver)
222 }
223 };
224
225 if let Some(release) = release {
226 let _ = release.await;
227 }
228 }
229
230 pub async fn reached(&self) {
232 let notification = {
233 let mut state = self.state.borrow_mut();
234 if state.reached > 0 {
235 None
236 } else {
237 let (sender, receiver) = oneshot::channel();
238 state.reached_notifiers.push(sender);
239 Some(receiver)
240 }
241 };
242
243 if let Some(notification) = notification {
244 let _ = notification.await;
245 }
246 }
247
248 #[must_use]
252 pub fn release(&self) -> bool {
253 let notifiers = {
254 let mut state = self.state.borrow_mut();
255 if state.released {
256 return false;
257 }
258 state.released = true;
259 std::mem::take(&mut state.release_notifiers)
260 };
261 for notifier in notifiers {
262 let _ = notifier.send(());
263 }
264 true
265 }
266
267 pub fn reached_count(&self) -> usize {
269 self.state.borrow().reached
270 }
271
272 pub fn is_released(&self) -> bool {
274 self.state.borrow().released
275 }
276}
277
278#[cfg(test)]
279mod tests {
280 use std::{cell::RefCell, rc::Rc};
281
282 use lenso_kernel::TaskOutcome;
283
284 use super::*;
285
286 #[test]
287 fn gate_blocks_work_until_the_test_releases_it() {
288 let simulator = TestSimulator::new();
289 let gate = simulator.gate("provider.before-response");
290 let task = simulator
291 .spawn({
292 let gate = gate.clone();
293 async move {
294 gate.wait().await;
295 }
296 })
297 .unwrap();
298
299 simulator.run(gate.reached());
300 assert_eq!(gate.name(), "provider.before-response");
301 assert_eq!(gate.reached_count(), 1);
302 assert!(!gate.is_released());
303 assert!(gate.release());
304 assert!(!gate.release());
305 assert_eq!(simulator.run(task), TaskOutcome::Completed);
306 }
307
308 #[test]
309 fn gates_allow_test_chosen_completion_order() {
310 let simulator = TestSimulator::new();
311 let first = simulator.gate("operation.first");
312 let second = simulator.gate("operation.second");
313 let completed = Rc::new(RefCell::new(Vec::new()));
314
315 let first_task = simulator
316 .spawn({
317 let first = first.clone();
318 let completed = completed.clone();
319 async move {
320 first.wait().await;
321 completed.borrow_mut().push("first");
322 }
323 })
324 .unwrap();
325 let second_task = simulator
326 .spawn({
327 let second = second.clone();
328 let completed = completed.clone();
329 async move {
330 second.wait().await;
331 completed.borrow_mut().push("second");
332 }
333 })
334 .unwrap();
335
336 simulator.run(async {
337 first.reached().await;
338 second.reached().await;
339 });
340 assert!(second.release());
341 assert_eq!(simulator.run(second_task), TaskOutcome::Completed);
342 assert_eq!(&*completed.borrow(), &["second"]);
343
344 assert!(first.release());
345 assert_eq!(simulator.run(first_task), TaskOutcome::Completed);
346 assert_eq!(&*completed.borrow(), &["second", "first"]);
347 }
348
349 #[test]
350 fn cancelling_one_blocked_operation_does_not_release_its_peer() {
351 let simulator = TestSimulator::new();
352 let cancelled = simulator.gate("operation.cancelled");
353 let peer = simulator.gate("operation.peer");
354 let cancelled_task = simulator
355 .spawn({
356 let cancelled = cancelled.clone();
357 async move {
358 cancelled.wait().await;
359 }
360 })
361 .unwrap();
362 let peer_task = simulator
363 .spawn({
364 let peer = peer.clone();
365 async move {
366 peer.wait().await;
367 }
368 })
369 .unwrap();
370
371 simulator.run(async {
372 cancelled.reached().await;
373 peer.reached().await;
374 });
375 cancelled_task.cancel();
376 assert_eq!(simulator.run(cancelled_task), TaskOutcome::Cancelled);
377 assert!(!peer.is_released());
378
379 assert!(peer.release());
380 assert_eq!(simulator.run(peer_task), TaskOutcome::Completed);
381 }
382
383 #[test]
384 fn advancing_virtual_time_wakes_elapsed_work_when_the_test_pumps() {
385 let simulator = TestSimulator::new();
386 let completed = Rc::new(RefCell::new(false));
387 let driver = simulator.driver();
388 let task = simulator
389 .spawn({
390 let completed = completed.clone();
391 async move {
392 driver.sleep_until(Duration::from_millis(5)).await;
393 *completed.borrow_mut() = true;
394 }
395 })
396 .unwrap();
397
398 simulator.pump();
399 assert!(!*completed.borrow());
400 simulator.advance(Duration::from_millis(5));
401 simulator.pump();
402 assert!(*completed.borrow());
403 assert_eq!(simulator.run(task), TaskOutcome::Completed);
404 }
405
406 #[test]
407 fn freezing_one_resource_owner_does_not_block_another() {
408 let simulator = TestSimulator::new();
409 let auth_store = simulator.resource("auth-store");
410 let object_store = simulator.resource("object-store");
411 assert!(auth_store.freeze());
412 assert!(!auth_store.freeze());
413
414 let blocked = simulator
415 .spawn({
416 let auth_store = auth_store.clone();
417 async move {
418 auth_store.acquire().await;
419 }
420 })
421 .unwrap();
422 let independent = simulator
423 .spawn({
424 let object_store = object_store.clone();
425 async move {
426 object_store.acquire().await;
427 }
428 })
429 .unwrap();
430
431 simulator.pump();
432 assert_eq!(auth_store.acquisition_count(), 0);
433 assert_eq!(object_store.acquisition_count(), 1);
434 assert_eq!(simulator.run(independent), TaskOutcome::Completed);
435 assert!(auth_store.is_frozen());
436
437 assert!(auth_store.thaw());
438 assert!(!auth_store.thaw());
439 assert_eq!(simulator.run(blocked), TaskOutcome::Completed);
440 assert_eq!(auth_store.acquisition_count(), 1);
441 }
442}