watermelon_testkit/
time.rs

1use std::{
2    future::{Future, IntoFuture},
3    panic::Location,
4    time::Duration,
5};
6
7use tokio::time::Instant;
8
9/// Bound how long a test is willing to wait on a future
10pub trait TimeoutExt: IntoFuture + Sized {
11    /// Await `self`, panicking if it doesn't complete within `timeout`
12    ///
13    /// The panic message points at the caller, so that a hung test
14    /// reports what it was waiting on instead of stalling CI.
15    #[track_caller]
16    fn within(self, timeout: Duration) -> impl Future<Output = Self::Output> {
17        let caller = Location::caller();
18        async move {
19            tokio::time::timeout(timeout, self)
20                .await
21                .unwrap_or_else(|_elapsed| {
22                    panic!("future at {caller} did not complete within {timeout:?}")
23                })
24        }
25    }
26}
27
28impl<F: IntoFuture> TimeoutExt for F {}
29
30/// Poll `condition` until it returns `true`, panicking after `timeout`
31///
32/// Only use this for state the test cannot synchronize with directly,
33/// like the server noticing that a TCP connection was closed.
34///
35/// # Panics
36///
37/// Panics if `condition` is still `false` after `timeout`.
38#[track_caller]
39pub fn eventually(
40    timeout: Duration,
41    mut condition: impl AsyncFnMut() -> bool,
42) -> impl Future<Output = ()> {
43    let caller = Location::caller();
44    async move {
45        let deadline = Instant::now() + timeout;
46        while !condition().await {
47            assert!(
48                Instant::now() < deadline,
49                "condition at {caller} did not become true within {timeout:?}"
50            );
51            tokio::time::sleep(Duration::from_millis(10)).await;
52        }
53    }
54}
watermelon_testkit/time.rs

watermelon_testkit/
time.rs
