tokio/tokio-timer/tests/throttle.rs

#![deny(warnings, rust_2018_idioms)]
#![cfg(feature = "async-traits")]

use tokio_sync::mpsc;
use tokio_test::task::MockTask;
use tokio_test::{assert_pending, assert_ready_eq, clock};
use tokio_timer::throttle::Throttle;

use futures_core::Stream;
use std::time::Duration;

macro_rules! poll {
    ($task:ident, $stream:ident) => {{
        use std::pin::Pin;
        $task.enter(|cx| Pin::new(&mut $stream).poll_next(cx))
    }};
}

#[test]
fn throttle() {
    let mut t = MockTask::new();

    clock::mock(|clock| {
        let (mut tx, rx) = mpsc::unbounded_channel();
        let mut stream = Throttle::new(rx, ms(1));

        assert_pending!(poll!(t, stream));

        for i in 0..3 {
            tx.try_send(i).unwrap();
        }

        for i in 0..3 {
            assert_ready_eq!(poll!(t, stream), Some(i));
            assert_pending!(poll!(t, stream));

            clock.advance(ms(1));
        }

        assert_pending!(poll!(t, stream));
    });
}

#[test]
fn throttle_dur_0() {
    let mut t = MockTask::new();

    clock::mock(|_| {
        let (mut tx, rx) = mpsc::unbounded_channel();
        let mut stream = Throttle::new(rx, ms(0));

        assert_pending!(poll!(t, stream));

        for i in 0..3 {
            tx.try_send(i).unwrap();
        }

        for i in 0..3 {
            assert_ready_eq!(poll!(t, stream), Some(i));
        }

        assert_pending!(poll!(t, stream));
    });
}

fn ms(n: u64) -> Duration {
    Duration::from_millis(n)
}