1 files changed, 58 insertions, 102 deletions

diff --git a/tokio-sync/tests/oneshot.rs b/tokio-sync/tests/oneshot.rs
index cc14dc95..fc78bd2e 100644
--- a/tokio-sync/tests/oneshot.rs
+++ b/tokio-sync/tests/oneshot.rs
@@ -1,29 +1,8 @@
 #![deny(warnings, rust_2018_idioms)]
 
-use futures;
-use futures::prelude::*;
-use tokio_mock_task::*;
 use tokio_sync::oneshot;
-
-macro_rules! assert_ready {
-    ($e:expr) => {{
-        match $e {
-            Ok(futures::Async::Ready(v)) => v,
-            Ok(_) => panic!("not ready"),
-            Err(e) => panic!("error = {:?}", e),
-        }
-    }};
-}
-
-macro_rules! assert_not_ready {
-    ($e:expr) => {{
-        match $e {
-            Ok(futures::Async::NotReady) => {}
-            Ok(futures::Async::Ready(v)) => panic!("ready; value = {:?}", v),
-            Err(e) => panic!("error = {:?}", e),
-        }
-    }};
-}
+use tokio_test::task::MockTask;
+use tokio_test::*;
 
 trait AssertSend: Send {}
 impl AssertSend for oneshot::Sender<i32> {}
@@ -34,15 +13,13 @@ fn send_recv() {
     let (tx, mut rx) = oneshot::channel();
     let mut task = MockTask::new();
 
-    task.enter(|| {
-        assert_not_ready!(rx.poll());
-    });
+    assert_pending!(task.poll(&mut rx));
 
-    assert!(tx.send(1).is_ok());
+    assert_ok!(tx.send(1));
 
-    assert!(task.is_notified());
+    assert!(task.is_woken());
 
-    let val = assert_ready!(rx.poll());
+    let val = assert_ready_ok!(task.poll(&mut rx));
     assert_eq!(val, 1);
 }
 
@@ -51,14 +28,12 @@ fn close_tx() {
     let (tx, mut rx) = oneshot::channel::<i32>();
     let mut task = MockTask::new();
 
-    task.enter(|| {
-        assert_not_ready!(rx.poll());
-    });
+    assert_pending!(task.poll(&mut rx));
 
     drop(tx);
 
-    assert!(task.is_notified());
-    assert!(rx.poll().is_err());
+    assert!(task.is_woken());
+    assert_ready_err!(task.poll(&mut rx));
 }
 
 #[test]
@@ -67,67 +42,64 @@ fn close_rx() {
     //
     let (tx, _) = oneshot::channel();
 
-    assert!(tx.send(1).is_err());
+    assert_err!(tx.send(1));
 
     // Second, via poll_close();
 
     let (mut tx, rx) = oneshot::channel();
     let mut task = MockTask::new();
 
-    task.enter(|| assert_not_ready!(tx.poll_close()));
+    assert_pending!(task.enter(|cx| tx.poll_close(cx)));
 
     drop(rx);
 
-    assert!(task.is_notified());
+    assert!(task.is_woken());
     assert!(tx.is_closed());
-    assert_ready!(tx.poll_close());
+    assert_ready!(task.enter(|cx| tx.poll_close(cx)));
 
-    assert!(tx.send(1).is_err());
+    assert_err!(tx.send(1));
 }
 
 #[test]
 fn explicit_close_poll() {
     // First, with message sent
     let (tx, mut rx) = oneshot::channel();
+    let mut task = MockTask::new();
 
-    assert!(tx.send(1).is_ok());
+    assert_ok!(tx.send(1));
 
     rx.close();
 
-    let value = assert_ready!(rx.poll());
+    let value = assert_ready_ok!(task.poll(&mut rx));
     assert_eq!(value, 1);
 
-    println!("~~~~~~~~~ TWO ~~~~~~~~~~");
-
     // Second, without the message sent
     let (mut tx, mut rx) = oneshot::channel::<i32>();
-    let mut task = MockTask::new();
 
-    task.enter(|| assert_not_ready!(tx.poll_close()));
+    assert_pending!(task.enter(|cx| tx.poll_close(cx)));
 
     rx.close();
 
-    assert!(task.is_notified());
+    assert!(task.is_woken());
     assert!(tx.is_closed());
-    assert_ready!(tx.poll_close());
-
-    assert!(tx.send(1).is_err());
+    assert_ready!(task.enter(|cx| tx.poll_close(cx)));
 
-    assert!(rx.poll().is_err());
+    assert_err!(tx.send(1));
+    assert_ready_err!(task.poll(&mut rx));
 
     // Again, but without sending the value this time
     let (mut tx, mut rx) = oneshot::channel::<i32>();
     let mut task = MockTask::new();
 
-    task.enter(|| assert_not_ready!(tx.poll_close()));
+    assert_pending!(task.enter(|cx| tx.poll_close(cx)));
 
     rx.close();
 
-    assert!(task.is_notified());
+    assert!(task.is_woken());
     assert!(tx.is_closed());
-    assert_ready!(tx.poll_close());
+    assert_ready!(task.enter(|cx| tx.poll_close(cx)));
 
-    assert!(rx.poll().is_err());
+    assert_ready_err!(task.poll(&mut rx));
 }
 
 #[test]
@@ -135,27 +107,26 @@ fn explicit_close_try_recv() {
     // First, with message sent
     let (tx, mut rx) = oneshot::channel();
 
-    assert!(tx.send(1).is_ok());
+    assert_ok!(tx.send(1));
 
     rx.close();
 
-    assert_eq!(rx.try_recv().unwrap(), 1);
-
-    println!("~~~~~~~~~ TWO ~~~~~~~~~~");
+    let val = assert_ok!(rx.try_recv());
+    assert_eq!(1, val);
 
     // Second, without the message sent
     let (mut tx, mut rx) = oneshot::channel::<i32>();
     let mut task = MockTask::new();
 
-    task.enter(|| assert_not_ready!(tx.poll_close()));
+    assert_pending!(task.enter(|cx| tx.poll_close(cx)));
 
     rx.close();
 
-    assert!(task.is_notified());
+    assert!(task.is_woken());
     assert!(tx.is_closed());
-    assert_ready!(tx.poll_close());
+    assert_ready!(task.enter(|cx| tx.poll_close(cx)));
 
-    assert!(rx.try_recv().is_err());
+    assert_err!(rx.try_recv());
 }
 
 #[test]
@@ -166,11 +137,9 @@ fn close_try_recv_poll() {
 
     rx.close();
 
-    assert!(rx.try_recv().is_err());
+    assert_err!(rx.try_recv());
 
-    task.enter(|| {
-        let _ = rx.poll();
-    });
+    let _ = task.poll(&mut rx);
 }
 
 #[test]
@@ -179,19 +148,14 @@ fn drops_tasks() {
     let mut tx_task = MockTask::new();
     let mut rx_task = MockTask::new();
 
-    tx_task.enter(|| {
-        assert_not_ready!(tx.poll_close());
-    });
-
-    rx_task.enter(|| {
-        assert_not_ready!(rx.poll());
-    });
+    assert_pending!(tx_task.enter(|cx| tx.poll_close(cx)));
+    assert_pending!(rx_task.poll(&mut rx));
 
     drop(tx);
     drop(rx);
 
-    assert_eq!(1, tx_task.notifier_ref_count());
-    assert_eq!(1, rx_task.notifier_ref_count());
+    assert_eq!(1, tx_task.waker_ref_count());
+    assert_eq!(1, rx_task.waker_ref_count());
 }
 
 #[test]
@@ -201,26 +165,22 @@ fn receiver_changes_task() {
     let mut task1 = MockTask::new();
     let mut task2 = MockTask::new();
 
-    task1.enter(|| {
-        assert_not_ready!(rx.poll());
-    });
+    assert_pending!(task1.poll(&mut rx));
 
-    assert_eq!(2, task1.notifier_ref_count());
-    assert_eq!(1, task2.notifier_ref_count());
+    assert_eq!(2, task1.waker_ref_count());
+    assert_eq!(1, task2.waker_ref_count());
 
-    task2.enter(|| {
-        assert_not_ready!(rx.poll());
-    });
+    assert_pending!(task2.poll(&mut rx));
 
-    assert_eq!(1, task1.notifier_ref_count());
-    assert_eq!(2, task2.notifier_ref_count());
+    assert_eq!(1, task1.waker_ref_count());
+    assert_eq!(2, task2.waker_ref_count());
 
-    tx.send(1).unwrap();
+    assert_ok!(tx.send(1));
 
-    assert!(!task1.is_notified());
-    assert!(task2.is_notified());
+    assert!(!task1.is_woken());
+    assert!(task2.is_woken());
 
-    assert_ready!(rx.poll());
+    assert_ready_ok!(task2.poll(&mut rx));
 }
 
 #[test]
@@ -230,24 +190,20 @@ fn sender_changes_task() {
     let mut task1 = MockTask::new();
     let mut task2 = MockTask::new();
 
-    task1.enter(|| {
-        assert_not_ready!(tx.poll_close());
-    });
+    assert_pending!(task1.enter(|cx| tx.poll_close(cx)));
 
-    assert_eq!(2, task1.notifier_ref_count());
-    assert_eq!(1, task2.notifier_ref_count());
+    assert_eq!(2, task1.waker_ref_count());
+    assert_eq!(1, task2.waker_ref_count());
 
-    task2.enter(|| {
-        assert_not_ready!(tx.poll_close());
-    });
+    assert_pending!(task2.enter(|cx| tx.poll_close(cx)));
 
-    assert_eq!(1, task1.notifier_ref_count());
-    assert_eq!(2, task2.notifier_ref_count());
+    assert_eq!(1, task1.waker_ref_count());
+    assert_eq!(2, task2.waker_ref_count());
 
     drop(rx);
 
-    assert!(!task1.is_notified());
-    assert!(task2.is_notified());
+    assert!(!task1.is_woken());
+    assert!(task2.is_woken());
 
-    assert_ready!(tx.poll_close());
+    assert_ready!(task2.enter(|cx| tx.poll_close(cx)));
 }