Repository: MystenLabs/sui
File: crates/sui-indexer-alt-framework/src/ingestion/streaming_client.rs (Rust)

Code before:

    enum StreamAction {
        Checkpoint(u64),
        Error,
        Timeout {
            deadline: Option<Instant>,
            duration: Duration,
        },
    }

    struct MockStreamState {
        actions: Arc<Mutex<Vec<StreamAction>>>,
    }

    impl Stream for MockStreamState {
        type Item = Result<Checkpoint>;

        fn poll_next(
            self: Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<Option<Self::Item>> {
            let mut actions = self.actions.lock().unwrap();
            if actions.is_empty() {
                return std::task::Poll::Ready(None);
            }
            let action = actions.remove(0);

            match action {
                StreamAction::Checkpoint(seq) => {
                    let mut builder = TestCheckpointBuilder::new(seq);
                    std::task::Poll::Ready(Some(Ok(builder.build_checkpoint())))
                }
                StreamAction::Error => std::task::Poll::Ready(Some(Err(Error::StreamingError(
                    anyhow::anyhow!("Mock streaming error"),
                )))),
                StreamAction::Timeout { deadline, duration } => {
                    match deadline {
                        None => {
                            // First poll: set deadline and re-insert
                            let deadline = Instant::now() + duration;
                            actions.insert(
                                0,
                                StreamAction::Timeout {
                                    deadline: Some(deadline),
                                    duration,
                                },
                            );
                            std::task::Poll::Pending
                        }
                        Some(deadline_instant) => {
                            if Instant::now() >= deadline_instant {

Review comment (suggestion):
It would be easier to implement this by peeking the first action, and then call `actions.remove(0)` only in the cases where you have actually consumed the action, but otherwise, for timeouts, just modify it in place.

Code after:
        Ok(Box::pin(converted_stream))
    }
}

#[cfg(test)]
pub mod test_utils {
    use super::*;
    use crate::types::test_checkpoint_data_builder::TestCheckpointBuilder;
    use std::sync::{Arc, Mutex};
    use std::time::{Duration, Instant};

    enum StreamAction {
        Checkpoint(u64),
        Error,
        Timeout {
            deadline: Option<Instant>,
            duration: Duration,
        },
    }

    struct MockStreamState {
        actions: Arc<Mutex<Vec<StreamAction>>>,
    }

    impl Stream for MockStreamState {
        type Item = Result<Checkpoint>;

        fn poll_next(
            self: Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<Option<Self::Item>> {
            let mut actions = self.actions.lock().unwrap();
            if actions.is_empty() {
                return std::task::Poll::Ready(None);
            }

            match &actions[0] {
                StreamAction::Checkpoint(seq) => {
                    let seq = *seq;
                    actions.remove(0);
                    let mut builder = TestCheckpointBuilder::new(seq);
                    std::task::Poll::Ready(Some(Ok(builder.build_checkpoint())))
                }
                StreamAction::Error => {
                    actions.remove(0);
                    std::task::Poll::Ready(Some(Err(Error::StreamingError(anyhow::anyhow!(
                        "Mock streaming error"
                    )))))
                }
                StreamAction::Timeout { deadline, duration } => match deadline {
                    None => {