From 2a30e13f38b864807f9ad92023e91b060a6227a4 Mon Sep 17 00:00:00 2001
From: cssivision <cssivision@gmail.com>
Date: Thu, 10 Dec 2020 15:36:43 +0800
Subject: net: expose poll_* methods on UnixDatagram (#3223)

---
 tokio/tests/uds_datagram.rs | 96 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 96 insertions(+)

(limited to 'tokio/tests/uds_datagram.rs')

diff --git a/tokio/tests/uds_datagram.rs b/tokio/tests/uds_datagram.rs
index ec2f6f82..cdabd7b1 100644
--- a/tokio/tests/uds_datagram.rs
+++ b/tokio/tests/uds_datagram.rs
@@ -2,6 +2,8 @@
 #![cfg(feature = "full")]
 #![cfg(unix)]
 
+use futures::future::poll_fn;
+use tokio::io::ReadBuf;
 use tokio::net::UnixDatagram;
 use tokio::try_join;
 
@@ -82,6 +84,8 @@ async fn try_send_recv_never_block() -> io::Result<()> {
 
     // Send until we hit the OS `net.unix.max_dgram_qlen`.
     loop {
+        dgram1.writable().await.unwrap();
+
         match dgram1.try_send(payload) {
             Err(err) => match err.kind() {
                 io::ErrorKind::WouldBlock | io::ErrorKind::Other => break,
@@ -96,6 +100,7 @@ async fn try_send_recv_never_block() -> io::Result<()> {
 
     // Read every dgram we sent.
     while count > 0 {
+        dgram2.readable().await.unwrap();
         let len = dgram2.try_recv(&mut recv_buf[..])?;
         assert_eq!(len, payload.len());
         assert_eq!(payload, &recv_buf[..len]);
@@ -134,3 +139,94 @@ async fn split() -> std::io::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn send_to_recv_from_poll() -> std::io::Result<()> {
+    let dir = tempfile::tempdir().unwrap();
+    let sender_path = dir.path().join("sender.sock");
+    let receiver_path = dir.path().join("receiver.sock");
+
+    let sender = UnixDatagram::bind(&sender_path)?;
+    let receiver = UnixDatagram::bind(&receiver_path)?;
+
+    let msg = b"hello";
+    poll_fn(|cx| sender.poll_send_to(cx, msg, &receiver_path)).await?;
+
+    let mut recv_buf = [0u8; 32];
+    let mut read = ReadBuf::new(&mut recv_buf);
+    let addr = poll_fn(|cx| receiver.poll_recv_from(cx, &mut read)).await?;
+
+    assert_eq!(read.filled(), msg);
+    assert_eq!(addr.as_pathname(), Some(sender_path.as_ref()));
+    Ok(())
+}
+
+#[tokio::test]
+async fn send_recv_poll() -> std::io::Result<()> {
+    let dir = tempfile::tempdir().unwrap();
+    let sender_path = dir.path().join("sender.sock");
+    let receiver_path = dir.path().join("receiver.sock");
+
+    let sender = UnixDatagram::bind(&sender_path)?;
+    let receiver = UnixDatagram::bind(&receiver_path)?;
+
+    sender.connect(&receiver_path)?;
+    receiver.connect(&sender_path)?;
+
+    let msg = b"hello";
+    poll_fn(|cx| sender.poll_send(cx, msg)).await?;
+
+    let mut recv_buf = [0u8; 32];
+    let mut read = ReadBuf::new(&mut recv_buf);
+    let _len = poll_fn(|cx| receiver.poll_recv(cx, &mut read)).await?;
+
+    assert_eq!(read.filled(), msg);
+    Ok(())
+}
+
+#[tokio::test]
+async fn try_send_to_recv_from() -> std::io::Result<()> {
+    let dir = tempfile::tempdir().unwrap();
+    let server_path = dir.path().join("server.sock");
+    let client_path = dir.path().join("client.sock");
+
+    // Create listener
+    let server = UnixDatagram::bind(&server_path)?;
+
+    // Create socket pair
+    let client = UnixDatagram::bind(&client_path)?;
+
+    for _ in 0..5 {
+        loop {
+            client.writable().await?;
+
+            match client.try_send_to(b"hello world", &server_path) {
+                Ok(n) => {
+                    assert_eq!(n, 11);
+                    break;
+                }
+                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
+                Err(e) => panic!("{:?}", e),
+            }
+        }
+
+        loop {
+            server.readable().await?;
+
+            let mut buf = [0; 512];
+
+            match server.try_recv_from(&mut buf) {
+                Ok((n, addr)) => {
+                    assert_eq!(n, 11);
+                    assert_eq!(addr.as_pathname(), Some(client_path.as_ref()));
+                    assert_eq!(&buf[0..11], &b"hello world"[..]);
+                    break;
+                }
+                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
+                Err(e) => panic!("{:?}", e),
+            }
+        }
+    }
+
+    Ok(())
+}
-- 
cgit v1.2.3