Testing futures and streams, how do I create a fake Context?

user_481812 0
Viewed 493

I'm trying to figure out how to unit test my Stream's fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {todo!()}

It takes a Context cx which is passed to nested stream. My tested and testing code can do without but I have to provide some to be able to test.

Rust playground:

use futures::ready;
use futures::stream::Stream;
use pin_project::pin_project;
use std::collections::VecDeque;
use std::pin::Pin;
use std::task::{Context, Poll}; // 0.4.22

#[pin_project]
pub struct MyStream<T>(#[pin] T);

impl<T: Stream> Stream for MyStream<T> {
    type Item = ();
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let _ignored = ready!(self.project().0.poll_next(cx));
        Poll::Ready(Some(()))
    }
}

#[test]
fn test1() {
    let setup = TestStream::from(vec![Poll::Ready(Some(true))]);
    let sut = MyStream(setup);

    // where do I get a fake Context<>?
    assert_eq!(Pin::new(&mut sut).poll_next(cx), Poll::Ready(Some(())))
}

#[pin_project]
struct TestStream<I> {
    items: VecDeque<Poll<Option<I>>>,
}
impl<T: IntoIterator<Item = Poll<Option<I>>>, I> From<T> for TestStream<I> {
    fn from(from: T) -> Self {
        TestStream {
            items: from.into_iter().collect(),
        }
    }
}
impl<I> Stream for TestStream<I> {
    type Item = I;
    fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if let Some(item) = self.project().items.pop_front() {
            item
        } else {
            Poll::Ready(None)
        }
    }
}
async-await rust future unit-testing
2 Answers

I've finally found the noop waker:

Playground:

#[test]
fn test1() {
    let waker = futures::task::noop_waker_ref();
    let mut cx = std::task::Context::from_waker(waker);
    let setup = TestStream::from(vec![Poll::Ready(Some(true))]);
    let mut sut = MyStream(setup);

    assert_eq!(Pin::new(&mut sut).poll_next(&mut cx), Poll::Ready(Some(())))
}

Thanks @rodrigo for the also correct answer.

user_4818120

Creating a Context is easy: Context::from_waker(waker: &'a Waker).

But, where do you get a fake Waker? You must use the unsafe Waker::from_raw(waker: RawWaker).

But... and the RawWaker? Well, there is this RawWaker::new(data: *const (), vtable: &'static RawWakerVTable). If you are not interested in actually running the futures, you can just use a static () pointer an a dummy vtable:

use std::task::{
    RawWaker,
    RawWakerVTable,
    Waker,
};

static DUMMY_VTABLE: RawWakerVTable = RawWakerVTable::new(
    dummy_clone,
    dummy_wake,
    dummy_wake_by_ref,
    dummy_drop,
);
unsafe fn dummy_clone(ptr: *const ()) -> RawWaker {
    RawWaker::new(
        ptr,
        &DUMMY_VTABLE
    )
}
unsafe fn dummy_wake(_ptr: *const ()) {
}

unsafe fn dummy_wake_by_ref(_ptr: *const ()) {
}

unsafe fn dummy_drop(_ptr: *const ()) {
}

fn dummy_waker() -> Waker {
    unsafe {
        Waker::from_raw(
            RawWaker::new(
                &(),
                &DUMMY_VTABLE
            )
        )
    }
}

Then you create a local waker and a context wherever you need them:

let waker = dummy_waker();
let mut cx = Context::new(&waker);
assert_eq!(Pin::new(&mut sut).poll_next(&mut cx), Poll::Ready(Some(())))
user_8658740
Related