diff --git a/tokio-test/src/io.rs b/tokio-test/src/io.rs index 4ec66a45d..2e16f1fe0 100644 --- a/tokio-test/src/io.rs +++ b/tokio-test/src/io.rs @@ -310,6 +310,8 @@ impl Inner { if now < until { break; + } else { + self.waiting = None; } } else { self.waiting = Some(Instant::now() + *dur); diff --git a/tokio-test/tests/io.rs b/tokio-test/tests/io.rs index f164abaf1..f182e37ee 100644 --- a/tokio-test/tests/io.rs +++ b/tokio-test/tests/io.rs @@ -2,6 +2,7 @@ use std::io; use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::time::{Duration, Instant}; use tokio_test::io::Builder; #[tokio::test] @@ -84,3 +85,65 @@ async fn mock_panics_write_data_left() { use tokio_test::io::Builder; Builder::new().write(b"write").build(); } + +#[tokio::test(start_paused = true)] +async fn wait() { + const FIRST_WAIT: Duration = Duration::from_secs(1); + + let mut mock = Builder::new() + .wait(FIRST_WAIT) + .read(b"hello ") + .read(b"world!") + .build(); + + let mut buf = [0; 256]; + + let start = Instant::now(); // record the time the read call takes + // + let n = mock.read(&mut buf).await.expect("read 1"); + assert_eq!(&buf[..n], b"hello "); + println!("time elapsed after first read {:?}", start.elapsed()); + + let n = mock.read(&mut buf).await.expect("read 2"); + assert_eq!(&buf[..n], b"world!"); + println!("time elapsed after second read {:?}", start.elapsed()); + + // make sure the .wait() instruction worked + assert!( + start.elapsed() >= FIRST_WAIT, + "consuming the whole mock only took {}ms", + start.elapsed().as_millis() + ); +} + +#[tokio::test(start_paused = true)] +async fn multiple_wait() { + const FIRST_WAIT: Duration = Duration::from_secs(1); + const SECOND_WAIT: Duration = Duration::from_secs(1); + + let mut mock = Builder::new() + .wait(FIRST_WAIT) + .read(b"hello ") + .wait(SECOND_WAIT) + .read(b"world!") + .build(); + + let mut buf = [0; 256]; + + let start = Instant::now(); // record the time it takes to consume the mock + + let n = mock.read(&mut buf).await.expect("read 1"); + assert_eq!(&buf[..n], b"hello "); + println!("time elapsed after first read {:?}", start.elapsed()); + + let n = mock.read(&mut buf).await.expect("read 2"); + assert_eq!(&buf[..n], b"world!"); + println!("time elapsed after second read {:?}", start.elapsed()); + + // make sure the .wait() instruction worked + assert!( + start.elapsed() >= FIRST_WAIT + SECOND_WAIT, + "consuming the whole mock only took {}ms", + start.elapsed().as_millis() + ); +}