use crate::pocket_ic::ClientError;
use crate::pocket_ic::PocketIc;
use std::{task::Poll, time::Duration};
#[derive(Debug, thiserror::Error)]
pub enum TickUntilError<E> {
#[error("PocketIC progression failed: {0}")]
Progress(#[source] ClientError),
#[error("readiness remained pending after {rounds} rounds")]
ProgressLimit {
rounds: u32,
},
#[error("readiness check failed: {0}")]
Failed(#[source] E),
}
pub fn advance_and_tick(pic: &PocketIc, advance: Duration, ticks: u32) -> Result<(), ClientError> {
if !advance.is_zero() {
pic.advance_time(advance)?;
}
for _ in 0..ticks {
pic.tick()?;
}
Ok(())
}
pub fn tick_until<T, E>(
pic: &PocketIc,
max_rounds: u32,
advance: Duration,
mut poll: impl FnMut(&PocketIc) -> Poll<Result<T, E>>,
) -> Result<T, TickUntilError<E>> {
poll_with_rounds(
max_rounds,
|| advance_and_tick(pic, advance, 1),
|| poll(pic),
)
}
fn poll_with_rounds<T, E>(
max_rounds: u32,
mut round: impl FnMut() -> Result<(), ClientError>,
mut poll: impl FnMut() -> Poll<Result<T, E>>,
) -> Result<T, TickUntilError<E>> {
let mut rounds = 0;
loop {
match poll() {
Poll::Ready(result) => return result.map_err(TickUntilError::Failed),
Poll::Pending if rounds == max_rounds => {
return Err(TickUntilError::ProgressLimit { rounds });
}
Poll::Pending => {
round().map_err(TickUntilError::Progress)?;
rounds += 1;
}
}
}
}
#[cfg(test)]
mod tests;