From b749013a54ee533ca035de40a8defb06a3ec6be5 Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Fri, 28 Oct 2022 13:55:13 -0400
Subject: [PATCH] test: Improve `tokio_test::task` docs (#5132)

---
 tokio-test/src/task.rs | 45 ++++++++++++++++++++++++++++++++++--------
 1 file changed, 37 insertions(+), 8 deletions(-)

diff --git a/tokio-test/src/task.rs b/tokio-test/src/task.rs
index f84b4f4c0..c1cfca162 100644
--- a/tokio-test/src/task.rs
+++ b/tokio-test/src/task.rs
@@ -1,4 +1,29 @@
-//! Futures task based helpers
+//! Futures task based helpers to easily test futures and manually written futures.
+//!
+//! The [`Spawn`] type is used as a mock task harness that allows you to poll futures
+//! without needing to setup pinning or context. Any future can be polled but if the
+//! future requires the tokio async context you will need to ensure that you poll the
+//! [`Spawn`] within a tokio context, this means that as long as you are inside the
+//! runtime it will work and you can poll it via [`Spawn`].
+//!
+//! [`Spawn`] also supports [`Stream`] to call `poll_next` without pinning
+//! or context.
+//!
+//! In addition to circumventing the need for pinning and context, [`Spawn`] also tracks
+//! the amount of times the future/task was woken. This can be useful to track if some
+//! leaf future notified the root task correctly.
+//!
+//! # Example
+//!
+//! ```
+//! use tokio_test::task;
+//!
+//! let fut = async {};
+//!
+//! let mut task = task::spawn(fut);
+//!
+//! assert!(task.poll().is_ready(), "Task was not ready!");
+//! ```
 
 #![allow(clippy::mutex_atomic)]
 
@@ -11,7 +36,11 @@ use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
 
 use tokio_stream::Stream;
 
-/// TODO: dox
+/// Spawn a future into a [`Spawn`] which wraps the future in a mocked executor.
+///
+/// This can be used to spawn a [`Future`] or a [`Stream`].
+///
+/// For more information, check the module docs.
 pub fn spawn<T>(task: T) -> Spawn<T> {
     Spawn {
         task: MockTask::new(),
@@ -19,16 +48,14 @@ pub fn spawn<T>(task: T) -> Spawn<T> {
     }
 }
 
-/// Future spawned on a mock task
+/// Future spawned on a mock task that can be used to poll the future or stream
+/// without needing pinning or context types.
 #[derive(Debug)]
 pub struct Spawn<T> {
     task: MockTask,
     future: Pin<Box<T>>,
 }
 
-/// Mock task
-///
-/// A mock task is able to intercept and track wake notifications.
 #[derive(Debug, Clone)]
 struct MockTask {
     waker: Arc<ThreadWaker>,
@@ -91,7 +118,8 @@ impl<T: Unpin> ops::DerefMut for Spawn<T> {
 }
 
 impl<T: Future> Spawn<T> {
-    /// Polls a future
+    /// If `T` is a [`Future`] then poll it. This will handle pinning and the context
+    /// type for the future.
     pub fn poll(&mut self) -> Poll<T::Output> {
         let fut = self.future.as_mut();
         self.task.enter(|cx| fut.poll(cx))
@@ -99,7 +127,8 @@ impl<T: Future> Spawn<T> {
 }
 
 impl<T: Stream> Spawn<T> {
-    /// Polls a stream
+    /// If `T` is a [`Stream`] then poll_next it. This will handle pinning and the context
+    /// type for the stream.
     pub fn poll_next(&mut self) -> Poll<Option<T::Item>> {
         let stream = self.future.as_mut();
         self.task.enter(|cx| stream.poll_next(cx))