67 lines
1.7 KiB
Rust
67 lines
1.7 KiB
Rust
use futures::FutureExt;
|
|
use std::future::Future;
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
|
|
/// Future extension helpers that are useful for tests
|
|
pub trait TestFuture: Future {
|
|
/// Drive `other` by polling `self`.
|
|
///
|
|
/// `self` must not resolve before `other` does.
|
|
fn drive<T>(&mut self, other: T) -> Drive<'_, Self, T>
|
|
where
|
|
T: Future,
|
|
Self: Future + Sized,
|
|
{
|
|
Drive {
|
|
driver: self,
|
|
future: Box::pin(other),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<T: Future> TestFuture for T {}
|
|
|
|
// ===== Drive ======
|
|
|
|
/// Drive a future to completion while also polling the driver
|
|
///
|
|
/// This is useful for H2 futures that also require the connection to be polled.
|
|
pub struct Drive<'a, T, U> {
|
|
driver: &'a mut T,
|
|
future: Pin<Box<U>>,
|
|
}
|
|
|
|
impl<'a, T, U> Future for Drive<'a, T, U>
|
|
where
|
|
T: Future + Unpin,
|
|
U: Future,
|
|
{
|
|
type Output = U::Output;
|
|
|
|
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
let mut looped = false;
|
|
loop {
|
|
match self.future.poll_unpin(cx) {
|
|
Poll::Ready(val) => return Poll::Ready(val),
|
|
Poll::Pending => {}
|
|
}
|
|
|
|
match self.driver.poll_unpin(cx) {
|
|
Poll::Ready(_) => {
|
|
if looped {
|
|
// Try polling the future one last time
|
|
panic!("driver resolved before future")
|
|
} else {
|
|
looped = true;
|
|
continue;
|
|
}
|
|
}
|
|
Poll::Pending => {}
|
|
}
|
|
|
|
return Poll::Pending;
|
|
}
|
|
}
|
|
}
|