Is it possible to make mocked AsyncWrite and AsyncRead behavior use tokio::time::sleep() and other async functions?

Thanks to lots of help from the community over in this question, I've managed to get a mocked test rig up and running that allows for arbitrary behavior when AsyncWrite and AsyncRead poll_* functions are called by the AsyncReadExt and AsyncWriteExt convenience wrappers. However, I never actually noticed until now that the functions of the aforementioned traits are not themselves marked async; a little digging suggests Rust doesn't fully support async trait functions yet? Regardless, this presents a problem as I was hoping to simulate I/O timeout with a call to tokio::time::sleep() inside my mocked poll_write() implementation, as follows:

trait AsyncReadWrite: AsyncRead + AsyncWrite + Unpin {}
mock! {
    AsyncReadWrite {}
    impl AsyncRead for AsyncReadWrite {
        fn poll_read<'ctx, 'buf>(
            self: Pin<&mut Self>,
            cx: &mut std::task::Context<'ctx>,
            buf: &mut ReadBuf<'buf>,
        ) -> Poll<Result<(), std::io::Error>>;
    }
    impl AsyncWrite for AsyncReadWrite {
        fn poll_write<'a>(
            self: Pin<&mut Self>,
            cx: &mut std::task::Context<'a>,
            buf: &[u8],
        ) -> Poll<Result<usize, std::io::Error>>;
        fn poll_flush<'a>(
            self: Pin<&mut Self>,
            cx: &mut std::task::Context<'a>,
        ) -> Poll<Result<(), std::io::Error>>;
        fn poll_shutdown<'a>(
            self: Pin<&mut Self>,
            cx: &mut std::task::Context<'a>,
        ) -> Poll<Result<(), std::io::Error>>;
    }
}

#[tokio::test]
async fn test_timeout() {
    struct MockNetworkerTimeout {}
    impl Networker for MockNetworkerTimeout {
        async fn connect<A: ToSocketAddrs>(
            &self,
            _addr: A,
        ) -> std::io::Result<impl AsyncRead + AsyncWrite + Unpin> {
            let mut mock_stream = MockAsyncReadWrite::new();
            mock_stream
                .expect_poll_read()
                .returning(|_, _| Poll::Ready(std::io::Result::Ok(())));
            mock_stream.expect_poll_write().returning(|ctx, buf: &[u8]| {
                dbg!("mock async writing timeout");
                sleep(Duration::from_secs(10)).await;
                Poll::Ready(Ok(1))
            });
            std::io::Result::Ok(mock_stream)
        }
    }
    let mock_networker = MockNetworkerTimeout {};

    match timeout(
            Duration::from_secs(5),
            // send_comms calls mock_networker.connect() and then calls AsyncWriteExt's write() on the resultant opaque AsyncRead + AsyncWrite + Unpin impl
            send_comms(&vec![0u8; 1024], &mock_networker)
        ).await {
            Ok(_) => panic!("send_comms succeeded despite delay of 10 seconds that should have caused timeout"),
            Err(e) => {
                dbg!(format!("Timed out as expected with {}", e));
            },
        }
}

But that fails to compile with the error:

error[E0728]: `await` is only allowed inside `async` functions and blocks
    --> ...
     |
1790 |                 mock_stream.expect_poll_write().returning(|ctx, buf: &[u8]| {
     |                                                           ----------------- this is not `async`
...
1796 |                     sleep(Duration::from_secs(10)).await;
     |                                                    ^^^^^ only allowed inside `async` functions and blocks

I can still simulate timeout by simply having the expect_poll_write() return Poll::Pending exclusively, but it would be nice to have a timing mechanism on the inside of the mocked behavior. Is there a way to get a handle to the expected async calling context inside these non-async trait functions?

That's a factor, but it's not the whole reason. poll_read() is not an async function because it doesn't return a Future, and this is an important element of its design. Calling it doesn't create a stateful object (future); it only tries to make some progress in writing, just like Future::poll() tries to make some progress towards completing the polled future.

ctx is the async context. You can hook it up to a Future by poll()ing the Future inside of poll_write(). (You must store the future β€” do not recreate it every poll! That will cause you to make no progress.)

1 Like

Hopefully little follow-up question: I'm having a bear of a time actually getting a Future to poll into my expect_poll_write() mock closure:

let sleep_pls = sleep(Duration::from_secs(10));
mock_stream.expect_poll_write().returning(|ctx, buf: &[u8]| {
    dbg!("mock async writing timeout");
    match sleep_pls.poll(ctx) {
        Poll::Ready(_) => Poll::Ready(Ok(1)),
        Poll::Pending => Poll::Pending,
    }
});

which gives me the strange error no method named poll found for struct Sleep in the current scope. I don't understand that since Sleep implements Future. The compiler suggests consider pinning the expression with std::pin::pin!() and assigning that to a new binding so I tried:

let sleep_pls = sleep(Duration::from_secs(10));
mock_stream.expect_poll_write().returning(|ctx, buf: &[u8]| {
    dbg!("mock async writing timeout");
    let pinned_future: Pin<&mut Sleep> = std::pin::pin!(sleep_pls);
    match pinned_future.poll(ctx) {
        Poll::Ready(_) => Poll::Ready(Ok(1)),
        Poll::Pending => Poll::Pending,
    }
});

the compiler complains that I can't move ownership of sleep_pls in the pin!() macro, but if I borrow it instead I get the same no method named poll... error for Pin<&mut &Sleep>; the compiler tells me poll() does exist on Pin<&mut Sleep>, but we can't move or clone the Sleep so... how does one make a Sleep to be polled later on?

Futures do need to be pinned to be polled. The brute force solution is Box::pin(sleep(...)) β€” that gives you a movable and pollable Pin<Box<Sleep>>. The efficient solution is to use pin_project on the struct that is to implement AsyncWrite to allow it to own the Sleep like an async block would.

Using pin projection would likely require you to not use your mocking framework, since it won’t cooperate. I would generally recommend that anyway β€” in my experience "expect exactly this call and return exactly this response" tends to result in tests that are fragile (need frequent updating when details of the circumstance change) and unrealistic (the mock is likely to be accidentally written to do something that a real implementation would never do). I personally recommend writing fakes (small and specialized, but correct, implementations of the relevant trait or protocol) instead of mocks whenever possible (unless the goal is to test for a precise sequence of events involving calls and returns where any change whatsoever is a bug). But you did not ask me for testing advice, and there is nothing wrong with using Box::pin in a test.

1 Like

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.