From 19a87e090ed528001e0363a30f6165304a710d49 Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Fri, 17 Apr 2020 15:37:59 -0400
Subject: [PATCH] test: Add `Future` and `Stream` impl for `Spawn`. (#2412)

---
 tokio-test/CHANGELOG.md |  4 ++++
 tokio-test/Cargo.toml   |  4 ++--
 tokio-test/src/lib.rs   |  2 +-
 tokio-test/src/task.rs  | 20 ++++++++++++++++++++
 4 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/tokio-test/CHANGELOG.md b/tokio-test/CHANGELOG.md
index e5a0093d5..50371c44b 100644
--- a/tokio-test/CHANGELOG.md
+++ b/tokio-test/CHANGELOG.md
@@ -1,3 +1,7 @@
+# 0.2.1 (April 17, 2020)
+
+- Add `Future` and `Stream` implementations for `task::Spawn<T>`.
+
 # 0.2.0 (November 25, 2019)
 
 - Initial release
diff --git a/tokio-test/Cargo.toml b/tokio-test/Cargo.toml
index a1e60500d..130035c29 100644
--- a/tokio-test/Cargo.toml
+++ b/tokio-test/Cargo.toml
@@ -7,13 +7,13 @@ name = "tokio-test"
 #   - Cargo.toml
 # - Update CHANGELOG.md.
 # - Create "v0.2.x" git tag.
-version = "0.2.0"
+version = "0.2.1"
 edition = "2018"
 authors = ["Tokio Contributors <team@tokio.rs>"]
 license = "MIT"
 repository = "https://github.com/tokio-rs/tokio"
 homepage = "https://tokio.rs"
-documentation = "https://docs.rs/tokio-test/0.2.0/tokio_test"
+documentation = "https://docs.rs/tokio-test/0.2.1/tokio_test"
 description = """
 Testing utilities for Tokio- and futures-based code
 """
diff --git a/tokio-test/src/lib.rs b/tokio-test/src/lib.rs
index c109c1481..185f317b9 100644
--- a/tokio-test/src/lib.rs
+++ b/tokio-test/src/lib.rs
@@ -1,4 +1,4 @@
-#![doc(html_root_url = "https://docs.rs/tokio-test/0.2.0")]
+#![doc(html_root_url = "https://docs.rs/tokio-test/0.2.1")]
 #![warn(
     missing_debug_implementations,
     missing_docs,
diff --git a/tokio-test/src/task.rs b/tokio-test/src/task.rs
index 04328e3d5..b31850baf 100644
--- a/tokio-test/src/task.rs
+++ b/tokio-test/src/task.rs
@@ -116,6 +116,26 @@ impl<T: Stream> Spawn<T> {
     }
 }
 
+impl<T: Future> Future for Spawn<T> {
+    type Output = T::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        // Safety: we only expose &mut T if T: Unpin therefore this is safe.
+        let future = unsafe { self.map_unchecked_mut(|s| &mut s.future) };
+        future.poll(cx)
+    }
+}
+
+impl<T: Stream> Stream for Spawn<T> {
+    type Item = T::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        // Safety: we only expose &mut T if T: Unpin therefore this is safe.
+        let stream = unsafe { self.map_unchecked_mut(|s| &mut s.future) };
+        stream.poll_next(cx)
+    }
+}
+
 impl MockTask {
     /// Creates new mock task
     fn new() -> Self {