From 7f7f74985e18a6bbe6e05e49ca2407506a000889 Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Mon, 15 Jul 2019 14:53:16 -0700
Subject: [PATCH] io: Minor adjustments to tokio-test IO (#1306)

This also re-exports `bytes::{Buf, BufMut}` from `tokio-io`.
---
 tokio-io/src/lib.rs  |  4 +++-
 tokio-test/src/io.rs | 12 +++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/tokio-io/src/lib.rs b/tokio-io/src/lib.rs
index ad8b833d9..d914f50c3 100644
--- a/tokio-io/src/lib.rs
+++ b/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};
diff --git a/tokio-test/src/io.rs b/tokio-test/src/io.rs
index 881ffe1a3..4e4ec1b3f 100644
--- a/tokio-test/src/io.rs
+++ b/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(()))
     }