From ccee8085bf06547c813ede004f35a32f0571f721 Mon Sep 17 00:00:00 2001
From: Audun Halland <audun.halland@pm.me>
Date: Mon, 28 Aug 2023 22:49:03 +0200
Subject: [PATCH] tmp: tokio mock

---
 Cargo.toml                  |  2 ++
 src/mock/mod.rs             |  1 +
 src/mock/tokio.rs           | 56 +++++++++++++++++++++++++++++++++++++
 tests/it/main.rs            |  3 ++
 tests/it/test_mock_tokio.rs | 34 ++++++++++++++++++++++
 5 files changed, 96 insertions(+)
 create mode 100644 src/mock/tokio.rs
 create mode 100644 tests/it/test_mock_tokio.rs

diff --git a/Cargo.toml b/Cargo.toml
index c4b56b9..7f45292 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,6 +19,7 @@ std = ["once_cell/std"]
 spin-lock = ["dep:spin"]
 mock-core = []
 mock-std = ["std", "mock-core"]
+mock-tokio = ["std", "dep:tokio"]
 nightly-tests = []
 unstable-doc-cfg = []
 critical-section = ["once_cell/critical-section"]
@@ -29,6 +30,7 @@ once_cell = { version = "1.17", default-features = false }
 polonius-the-crab = "0.3"
 pretty_assertions = { version = "1.3", optional = true }
 spin = { version = "0.9.8", optional = true }
+tokio = { version = "1", default-features = false, optional = true }
 
 [dev-dependencies]
 async-trait = "0.1"
diff --git a/src/mock/mod.rs b/src/mock/mod.rs
index cdaa2c9..d45918f 100644
--- a/src/mock/mod.rs
+++ b/src/mock/mod.rs
@@ -1,2 +1,3 @@
 pub mod core;
 pub mod std;
+pub mod tokio;
diff --git a/src/mock/tokio.rs b/src/mock/tokio.rs
new file mode 100644
index 0000000..e8cfb1c
--- /dev/null
+++ b/src/mock/tokio.rs
@@ -0,0 +1,56 @@
+//! Mock APIs for `tokio` traits
+
+/// Mock APIs for `tokio::io` traits
+#[cfg(feature = "mock-tokio")]
+pub mod io {
+    use core::pin::Pin;
+    use core::task::{Context, Poll};
+    use std::io::IoSlice;
+
+    use tokio::io::{ReadBuf, Result, SeekFrom};
+
+    use crate::unimock;
+
+    #[unimock(prefix=crate, api=AsyncBufReadMock, mirror=tokio::io::AsyncBufRead)]
+    pub trait AsyncBufRead {
+        fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<&[u8]>>;
+        fn consume(self: Pin<&mut Self>, amt: usize);
+    }
+
+    #[unimock(prefix=crate, api=AsyncReadMock, mirror=tokio::io::AsyncRead)]
+    pub trait AsyncRead {
+        fn poll_read(
+            self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+            buf: &mut ReadBuf<'_>,
+        ) -> Poll<Result<()>>;
+    }
+
+    #[unimock(prefix=crate, api=AsyncSeekMock, mirror=tokio::io::AsyncSeek)]
+    pub trait AsyncSeek {
+        fn start_seek(self: Pin<&mut Self>, position: SeekFrom) -> Result<()>;
+        fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<u64>>;
+    }
+
+    #[unimock(prefix=crate, api=AsyncWriteMock, mirror=tokio::io::AsyncWrite)]
+    pub trait AsyncWrite {
+        fn poll_write(
+            self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+            buf: &[u8],
+        ) -> Poll<Result<usize>>;
+
+        fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>>;
+
+        fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>>;
+
+        fn poll_write_vectored(
+            self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+            bufs: &[IoSlice<'_>],
+        ) -> Poll<Result<usize>> {
+        }
+
+        fn is_write_vectored(&self) -> bool {}
+    }
+}
diff --git a/tests/it/main.rs b/tests/it/main.rs
index 0d0db6b..d111446 100644
--- a/tests/it/main.rs
+++ b/tests/it/main.rs
@@ -31,6 +31,9 @@ mod async_fn;
 #[cfg(all(feature = "mock-core", feature = "mock-std"))]
 mod std;
 
+#[cfg(all(feature = "mock-tokio", feature = "std"))]
+mod test_mock_tokio;
+
 mod unmock;
 
 fn main() {}
diff --git a/tests/it/test_mock_tokio.rs b/tests/it/test_mock_tokio.rs
new file mode 100644
index 0000000..ecb8fa8
--- /dev/null
+++ b/tests/it/test_mock_tokio.rs
@@ -0,0 +1,34 @@
+use core::task::Poll;
+
+use unimock::{mock::tokio::io::AsyncReadMock, *};
+
+use tokio::io::AsyncReadExt;
+
+#[tokio::test]
+async fn test_tokio_read() {
+    let mut u = Unimock::new((
+        AsyncReadMock::poll_read
+            .next_call(matching!())
+            .mutates(|buf, (_, _)| {
+                buf.put_slice(&[1, 2, 3]);
+
+                // Can return Poll::Ready explicitly
+                Poll::Ready(Ok(()))
+            }),
+        AsyncReadMock::poll_read
+            .next_call(matching!())
+            .mutates(|buf, _| {
+                buf.put_slice(&[5, 6, 7]);
+
+                // Also can return just Ok(()) because of `impl From<T> for Poll<T>`
+                Ok(())
+            }),
+    ));
+    let mut buf = [0; 10];
+
+    let n = u.read(&mut buf).await.unwrap();
+    assert_eq!(n, 3);
+    assert_eq!(&buf[0..n], &[1, 2, 3]);
+
+    u.read(&mut buf).await.unwrap();
+}