From e53df802de52db8b1489a1a43166f441b12ca5c7 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Mon, 4 Nov 2019 12:26:40 -0800
Subject: [PATCH] chore: fix nightly ci (#416)

## Motivation

A recent upstream change to `tokio-test` (tokio-rs/tokio#1728) removed
the `MockTask` API in favour of a more ergonomic `task::spawn` API. This
broke our tests for `tracing-futures` and `tracing-attributes` `std::future`
compatibility, which depend on `tokio-test` from Git.

## Solution

This branch updates these tests to use the new API, which *is* much
nicer. This should fix nightly CI.

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
---
 tracing-attributes/test_async_await/tests/async_fn.rs | 4 +---
 tracing-futures/test_std_future/src/lib.rs            | 9 ++++-----
 tracing-futures/test_std_future/tests/test.rs         | 7 ++-----
 3 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/tracing-attributes/test_async_await/tests/async_fn.rs b/tracing-attributes/test_async_await/tests/async_fn.rs
index 38daadf330..f62ed34f95 100644
--- a/tracing-attributes/test_async_await/tests/async_fn.rs
+++ b/tracing-attributes/test_async_await/tests/async_fn.rs
@@ -6,7 +6,6 @@ use test_std_future::{
     support::*,
 };
 
-use tokio_test::task::MockTask;
 use tracing::subscriber::with_default;
 
 #[tracing_attributes::instrument]
@@ -28,9 +27,8 @@ fn async_fn_only_enters_for_polls() {
         .drop_span(span::mock().named("test_async_fn"))
         .done()
         .run_with_handle();
-    let mut task = MockTask::new();
     with_default(subscriber, || {
-        block_on_future(&mut task, async { test_async_fn(2).await }).unwrap();
+        block_on_future(async { test_async_fn(2).await }).unwrap();
     });
     handle.assert_finished();
 }
diff --git a/tracing-futures/test_std_future/src/lib.rs b/tracing-futures/test_std_future/src/lib.rs
index a3344553a8..6c1df69c3c 100644
--- a/tracing-futures/test_std_future/src/lib.rs
+++ b/tracing-futures/test_std_future/src/lib.rs
@@ -6,7 +6,7 @@ use std::{
     pin::Pin,
     task::{Context, Poll},
 };
-use tokio_test::task::MockTask;
+use tokio_test::task;
 
 pub struct PollN<T, E> {
     and_return: Option<Result<T, E>>,
@@ -53,14 +53,13 @@ impl PollN<(), ()> {
     }
 }
 
-pub fn block_on_future<F>(task: &mut MockTask, future: F) -> F::Output
+pub fn block_on_future<F>(future: F) -> F::Output
 where
     F: Future,
 {
-    let mut future = Box::pin(future);
-
+    let mut task = task::spawn(future);
     loop {
-        match task.poll(&mut future) {
+        match task.poll() {
             Poll::Ready(v) => break v,
             _ => {}
         }
diff --git a/tracing-futures/test_std_future/tests/test.rs b/tracing-futures/test_std_future/tests/test.rs
index 854b0a5305..673eca1879 100644
--- a/tracing-futures/test_std_future/tests/test.rs
+++ b/tracing-futures/test_std_future/tests/test.rs
@@ -4,7 +4,6 @@ extern crate tracing_core;
 
 use test_std_future::{block_on_future, support::*, PollN};
 
-use tokio_test::task::MockTask;
 use tracing::{subscriber::with_default, Level};
 use tracing_futures::Instrument;
 
@@ -18,10 +17,9 @@ fn enter_exit_is_reasonable() {
         .drop_span(span::mock().named("foo"))
         .done()
         .run_with_handle();
-    let mut task = MockTask::new();
     with_default(subscriber, || {
         let future = PollN::new_ok(2).instrument(span!(Level::TRACE, "foo"));
-        block_on_future(&mut task, future).unwrap();
+        block_on_future(future).unwrap();
     });
     handle.assert_finished();
 }
@@ -36,10 +34,9 @@ fn error_ends_span() {
         .drop_span(span::mock().named("foo"))
         .done()
         .run_with_handle();
-    let mut task = MockTask::new();
     with_default(subscriber, || {
         let future = PollN::new_err(2).instrument(span!(Level::TRACE, "foo"));
-        block_on_future(&mut task, future).unwrap_err();
+        block_on_future(future).unwrap_err();
     });
     handle.assert_finished();
 }