io: Minor adjustments to tokio-test IO (#1306)

This also re-exports `bytes::{Buf, BufMut}` from `tokio-io`.

tokio-io/src/lib.rs

@@ -2,7 +2,6 @@
 #![deny(missing_debug_implementations, missing_docs, rust_2018_idioms)]
 #![cfg_attr(test, deny(warnings))]
 #![doc(test(no_crate_inject, attr(deny(rust_2018_idioms))))]
-//#![feature(async_await)]
 
 //! Core I/O traits and combinators when working with Tokio.
 //!
@@ -19,3 +18,6 @@ mod async_write;
 pub use self::async_buf_read::AsyncBufRead;
 pub use self::async_read::AsyncRead;
 pub use self::async_write::AsyncWrite;
+
+// Re-export `Buf` and `BufMut` since they are part of the API
+pub use bytes::{Buf, BufMut};

tokio-test/src/io.rs

@@ -21,7 +21,7 @@ use std::time::{Duration, Instant};
 use std::{cmp, io};
 
 use futures_core::ready;
-use tokio_io::{AsyncRead, AsyncWrite};
+use tokio_io::{AsyncRead, AsyncWrite, Buf};
 use tokio_sync::mpsc;
 use tokio_timer::{clock, timer, Delay};
 
@@ -409,6 +409,16 @@ impl AsyncWrite for Mock {
         }
     }
 
+    fn poll_write_buf<B: Buf>(
+        self: Pin<&mut Self>,
+        cx: &mut task::Context<'_>,
+        buf: &mut B,
+    ) -> Poll<io::Result<usize>> {
+        let n = ready!(self.poll_write(cx, buf.bytes()))?;
+        buf.advance(n);
+        Poll::Ready(Ok(n))
+    }
+
     fn poll_flush(self: Pin<&mut Self>, _cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
         Poll::Ready(Ok(()))
     }