ic-testkit 0.46.0

use pocket_ic::PocketIc;
use std::{fmt, task::Poll, time::Duration};

/// A caller-owned readiness check failed or exhausted its simulated round budget.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum TickUntilError<E> {
    /// The predicate was still pending after this many advance/tick rounds.
    ProgressLimit { rounds: u32 },
    /// The predicate returned its own failure.
    Failed(E),
}

impl<E: fmt::Display> fmt::Display for TickUntilError<E> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::ProgressLimit { rounds } => {
                write!(
                    formatter,
                    "readiness remained pending after {rounds} rounds"
                )
            }
            Self::Failed(error) => write!(formatter, "readiness check failed: {error}"),
        }
    }
}

impl<E: std::error::Error + 'static> std::error::Error for TickUntilError<E> {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::ProgressLimit { .. } => None,
            Self::Failed(error) => Some(error),
        }
    }
}

/// Check a caller-owned predicate, advancing simulated time and ticking while pending.
///
/// The predicate runs once before any mutation. Each pending result permits one
/// round of `advance_time(advance)`, then `tick()`, then another predicate check,
/// up to `max_rounds`. A zero budget still performs the initial check. A zero
/// duration permits ticks without reading or writing simulated time. Completion
/// on the last allowed round succeeds; predicate failures return immediately without another round.
///
/// This bounds simulated progression, not wall-clock time spent inside PocketIC
/// calls or the predicate. PocketIC operation failures retain upstream panic
/// behavior. Applications own readiness meaning and predicate error types.
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,
        || {
            if !advance.is_zero() {
                pic.advance_time(advance);
            }
            pic.tick();
        },
        || poll(pic),
    )
}

fn poll_with_rounds<T, E>(
    max_rounds: u32,
    mut round: impl FnMut(),
    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();
                rounds += 1;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::cell::Cell;

    #[test]
    fn initial_completion_and_failure_do_not_advance() {
        for budget in [0, 3] {
            for result in [Ok(42), Err("predicate failure")] {
                assert_eq!(
                    poll_with_rounds(
                        budget,
                        || panic!("unexpected round"),
                        || Poll::Ready(result)
                    ),
                    result.map_err(TickUntilError::Failed)
                );
            }
        }
    }

    #[test]
    fn completion_on_last_round_preserves_poll_order() {
        let rounds = Cell::new(0);
        let polls = Cell::new(0);
        let result = poll_with_rounds(
            3,
            || {
                assert_eq!(polls.get(), rounds.get() + 1);
                rounds.set(rounds.get() + 1);
            },
            || {
                assert_eq!(polls.get(), rounds.get());
                polls.set(polls.get() + 1);
                if rounds.get() == 3 {
                    Poll::Ready(Ok::<_, ()>(42))
                } else {
                    Poll::Pending
                }
            },
        );
        assert_eq!(result, Ok(42));
        assert_eq!((rounds.get(), polls.get()), (3, 4));
    }

    #[test]
    fn exhaustion_performs_exactly_the_allowed_rounds() {
        for budget in [0, 1, 3] {
            let rounds = Cell::new(0);
            let polls = Cell::new(0);
            let result = poll_with_rounds(
                budget,
                || rounds.set(rounds.get() + 1),
                || {
                    polls.set(polls.get() + 1);
                    Poll::<Result<(), ()>>::Pending
                },
            );
            assert_eq!(
                result,
                Err(TickUntilError::ProgressLimit { rounds: budget })
            );
            assert_eq!((rounds.get(), polls.get()), (budget, budget + 1));
        }
    }

    #[test]
    fn failure_after_progress_preserves_the_original_cause() {
        let rounds = Cell::new(0);
        let result = poll_with_rounds(
            3,
            || rounds.set(rounds.get() + 1),
            || {
                if rounds.get() == 1 {
                    Poll::Ready(Err::<(), _>(std::io::Error::from(
                        std::io::ErrorKind::PermissionDenied,
                    )))
                } else {
                    Poll::Pending
                }
            },
        );
        let error = result.unwrap_err();
        let source = std::error::Error::source(&error).unwrap();
        assert_eq!(
            source.downcast_ref::<std::io::Error>().unwrap().kind(),
            std::io::ErrorKind::PermissionDenied
        );
        assert_eq!(rounds.get(), 1);
    }
    #[test]
    fn polling_transport_skips_zero_time_and_preserves_nonzero_advancement() {
        use std::{io::Write as _, net::TcpListener, thread, time::Instant};
        for advance in [Duration::ZERO, Duration::from_nanos(7)] {
            let listener = TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)).unwrap();
            listener.set_nonblocking(true).unwrap();
            let url = format!("http://{}/", listener.local_addr().unwrap())
                .parse()
                .unwrap();
            let worker = thread::spawn(move || {
                let deadline = Instant::now() + Duration::from_secs(5);
                let mut requests = Vec::new();
                loop {
                    assert!(Instant::now() < deadline, "polling peer timed out");
                    let (mut stream, _) = match listener.accept() {
                        Ok(connection) => connection,
                        Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                            thread::sleep(Duration::from_millis(1));
                            continue;
                        }
                        Err(error) => panic!("accept polling request: {error}"),
                    };
                    let request = crate::testing::read_http_request(
                        &mut stream,
                        Duration::from_secs(1),
                        16_384,
                        16_384,
                    )
                    .unwrap();
                    let text = std::str::from_utf8(&request).unwrap();
                    let path = text.split_whitespace().nth(1).unwrap().to_owned();
                    let reply = if path.ends_with("/read/get_time") {
                        "{\"nanos_since_epoch\":100}"
                    } else {
                        "null"
                    };
                    if path.ends_with("/update/set_time") {
                        let body = text.split_once("\r\n\r\n").unwrap().1;
                        assert_eq!(
                            serde_json::from_str::<serde_json::Value>(body).unwrap()["nanos_since_epoch"],
                            100 + u64::from(advance.subsec_nanos())
                        );
                    }
                    stream
                        .set_write_timeout(Some(Duration::from_secs(1)))
                        .unwrap();
                    write!(stream, "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{reply}", reply.len()).unwrap();
                    let finished = path.ends_with("/update/tick");
                    requests.push(path);
                    if finished {
                        return requests;
                    }
                }
            });
            let pic = PocketIc::new_from_existing_instance(url, 0, Some(1_000));
            let mut polls = 0;
            let result = tick_until(&pic, 1, advance, |_| {
                polls += 1;
                if polls == 2 {
                    Poll::Ready(Ok::<_, ()>(42))
                } else {
                    Poll::Pending
                }
            });
            assert_eq!(result, Ok(42));
            assert_eq!(polls, 2);
            let requests = worker.join().unwrap();
            let expected = if advance.is_zero() {
                vec!["/instances/0/update/tick"]
            } else {
                vec![
                    "/instances/0/read/get_time",
                    "/instances/0/update/set_time",
                    "/instances/0/update/tick",
                ]
            };
            assert_eq!(requests, expected);
        }
    }
}