Add expect_read test macro

Testing macro that calls read/recv/peek or recv_from/peek_from and checks if the received bytes are expected. In case of {recv,peek}_from it also checks the source address.
tokio-rs · Nov 22, 2019 · 9025e98 · 9025e98
1 parent 892f856
commit 9025e98
Show file tree

Hide file tree

Showing 6 changed files with 100 additions and 160 deletions.
diff --git a/tests/tcp.rs b/tests/tcp.rs
@@ -1,18 +1,18 @@
 use std::io::{self, Read, Write};
-use std::net::Shutdown;
+use std::net::{self, Shutdown};
 #[cfg(unix)]
 use std::os::unix::io::{FromRawFd, IntoRawFd};
 use std::sync::mpsc::channel;
-use std::thread::sleep;
+use std::thread::{self, sleep};
 use std::time::Duration;
-use std::{net, thread};
 
 #[cfg(unix)]
 use net2::TcpStreamExt;
 
 use mio::net::{TcpListener, TcpStream};
 use mio::{Events, Interests, Poll, Token};
 
+#[macro_use]
 mod util;
 
 use util::{
@@ -512,7 +512,7 @@ fn connection_reset_by_peer() {
     drop(client);
 
     // Wait a moment
-    thread::sleep(Duration::from_millis(100));
+    sleep(Duration::from_millis(100));
 
     // Register interest in the server socket
     poll.registry()
@@ -815,8 +815,7 @@ fn write_then_drop() {
     assert_eq!(events.iter().next().unwrap().token(), Token(3));
 
     let mut buf = [0; 10];
-    assert_eq!(s.read(&mut buf).unwrap(), 4);
-    assert_eq!(&buf[0..4], &[1, 2, 3, 4]);
+    expect_read!(s.read(&mut buf), &[1, 2, 3, 4]);
 }
 
 #[test]
@@ -871,8 +870,7 @@ fn write_then_deregister() {
     assert_eq!(events.iter().next().unwrap().token(), Token(3));
 
     let mut buf = [0; 10];
-    assert_eq!(s.read(&mut buf).unwrap(), 4);
-    assert_eq!(&buf[0..4], &[1, 2, 3, 4]);
+    expect_read!(s.read(&mut buf), &[1, 2, 3, 4]);
 }
 
 const ID1: Token = Token(1);
@@ -929,15 +927,13 @@ fn tcp_no_events_after_deregister() {
     expect_no_events(&mut poll, &mut events);
 
     let mut buf = [0; 10];
-    assert_eq!(stream.read(&mut buf).unwrap(), 4);
-    assert_eq!(&buf[0..4], &[1, 2, 3, 4]);
+    expect_read!(stream.read(&mut buf), &[1, 2, 3, 4]);
 
-    stream2.write_all(&[1, 2, 3, 4]).unwrap();
+    checked_write!(stream2.write(&[1, 2, 3, 4]));
     expect_no_events(&mut poll, &mut events);
 
     sleep(Duration::from_millis(200));
-    assert_eq!(stream.read(&mut buf).unwrap(), 4);
-    assert_eq!(&buf[0..4], &[1, 2, 3, 4]);
+    expect_read!(stream.read(&mut buf), &[1, 2, 3, 4]);
 
     expect_no_events(&mut poll, &mut events);
 }
diff --git a/tests/tcp_stream.rs b/tests/tcp_stream.rs
@@ -93,13 +93,8 @@ where
         vec![ExpectEvent::new(ID1, Interests::READABLE)],
     );
 
-    let n = stream.peek(&mut buf).expect("unable to peek from stream");
-    assert_eq!(n, DATA1.len());
-    assert_eq!(&buf[..n], DATA1);
-
-    let n = stream.read(&mut buf).expect("unable to read from stream");
-    assert_eq!(n, DATA1.len());
-    assert_eq!(&buf[..n], DATA1);
+    expect_read!(stream.peek(&mut buf), DATA1);
+    expect_read!(stream.read(&mut buf), DATA1);
 
     assert!(stream.take_error().unwrap().is_none());
 
@@ -168,9 +163,7 @@ fn try_clone() {
     );
 
     let mut buf = [0; 20];
-    let n = stream2.read(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(&buf[..n], DATA1);
+    expect_read!(stream2.read(&mut buf), DATA1);
 
     drop(stream2);
     thread_handle.join().expect("unable to join thread");
@@ -342,8 +335,7 @@ fn shutdown_read() {
     ))]
     {
         let mut buf = [0; 20];
-        let n = stream.read(&mut buf).unwrap();
-        assert_eq!(n, 0);
+        expect_read!(stream.read(&mut buf), &[]);
     }
 
     drop(stream);
@@ -385,9 +377,7 @@ fn shutdown_write() {
 
     // Read should be ok.
     let mut buf = [0; 20];
-    let n = stream.read(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(&buf[..n], DATA1);
+    expect_read!(stream.read(&mut buf), DATA1);
 
     drop(stream);
     thread_handle.join().expect("unable to join thread");
@@ -433,8 +423,7 @@ fn shutdown_both() {
     ))]
     {
         let mut buf = [0; 20];
-        let n = stream.read(&mut buf).unwrap();
-        assert_eq!(n, 0);
+        expect_read!(stream.read(&mut buf), &[]);
     }
 
     let err = stream.write(DATA2).unwrap_err();

diff --git a/tests/udp_socket.rs b/tests/udp_socket.rs
@@ -103,25 +103,11 @@ fn smoke_test_unconnected_udp_socket(socket1: UdpSocket, socket2: UdpSocket) {
         ],
     );
 
-    let (n, got_address1) = socket1.peek_from(&mut buf).unwrap();
-    assert_eq!(n, DATA2.len());
-    assert_eq!(buf[..n], DATA2[..]);
-    assert_eq!(got_address1, address2);
-
-    let (n, got_address2) = socket2.peek_from(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
-    assert_eq!(got_address2, address1);
-
-    let (n, got_address1) = socket1.recv_from(&mut buf).unwrap();
-    assert_eq!(n, DATA2.len());
-    assert_eq!(buf[..n], DATA2[..]);
-    assert_eq!(got_address1, address2);
-
-    let (n, got_address2) = socket2.recv_from(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
-    assert_eq!(got_address2, address1);
+    expect_read!(socket1.peek_from(&mut buf), DATA2, address2);
+    expect_read!(socket2.peek_from(&mut buf), DATA1, address1);
+
+    expect_read!(socket1.recv_from(&mut buf), DATA2, address2);
+    expect_read!(socket2.recv_from(&mut buf), DATA1, address1);
 
     assert!(socket1.take_error().unwrap().is_none());
     assert!(socket2.take_error().unwrap().is_none());
@@ -232,21 +218,11 @@ fn smoke_test_connected_udp_socket(socket1: UdpSocket, socket2: UdpSocket) {
     );
 
     let mut buf = [0; 20];
-    let n = socket1.peek(&mut buf).unwrap();
-    assert_eq!(n, DATA2.len());
-    assert_eq!(buf[..n], DATA2[..]);
-
-    let n = socket2.peek(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
-
-    let n = socket1.recv(&mut buf).unwrap();
-    assert_eq!(n, DATA2.len());
-    assert_eq!(buf[..n], DATA2[..]);
+    expect_read!(socket1.peek(&mut buf), DATA2);
+    expect_read!(socket2.peek(&mut buf), DATA1);
 
-    let n = socket2.recv(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
+    expect_read!(socket1.recv(&mut buf), DATA2);
+    expect_read!(socket2.recv(&mut buf), DATA1);
 
     assert!(socket1.take_error().unwrap().is_none());
     assert!(socket2.take_error().unwrap().is_none());
@@ -293,9 +269,7 @@ fn reconnect_udp_socket_sending() {
     );
 
     let mut buf = [0; 20];
-    let n = socket2.recv(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
+    expect_read!(socket2.recv(&mut buf), DATA1);
 
     socket1.connect(address3).unwrap();
     checked_write!(socket1.send(DATA2));
@@ -306,9 +280,7 @@ fn reconnect_udp_socket_sending() {
         vec![ExpectEvent::new(ID3, Interests::READABLE)],
     );
 
-    let n = socket3.recv(&mut buf).unwrap();
-    assert_eq!(n, DATA2.len());
-    assert_eq!(buf[..n], DATA2[..]);
+    expect_read!(socket3.recv(&mut buf), DATA2);
 
     assert!(socket1.take_error().unwrap().is_none());
     assert!(socket2.take_error().unwrap().is_none());
@@ -360,9 +332,7 @@ fn reconnect_udp_socket_receiving() {
     );
 
     let mut buf = [0; 20];
-    let n = socket1.recv(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
+    expect_read!(socket1.recv(&mut buf), DATA1);
 
     socket1.connect(address3).unwrap();
     checked_write!(socket3.send(DATA2));
@@ -375,9 +345,7 @@ fn reconnect_udp_socket_receiving() {
 
     // Read only a part of the data.
     let max = 4;
-    let n = socket1.recv(&mut buf[..max]).unwrap();
-    assert_eq!(n, max);
-    assert_eq!(buf[..max], DATA2[..max]);
+    expect_read!(socket1.recv(&mut buf[..max]), &DATA2[..max]);
 
     // Now connect back to socket 2, dropping the unread data.
     socket1.connect(address2).unwrap();
@@ -389,9 +357,7 @@ fn reconnect_udp_socket_receiving() {
         vec![ExpectEvent::new(ID1, Interests::READABLE)],
     );
 
-    let n = socket1.recv(&mut buf).unwrap();
-    assert_eq!(n, DATA2.len());
-    assert_eq!(buf[..n], DATA2[..]);
+    expect_read!(socket1.recv(&mut buf), DATA2);
 
     assert!(socket1.take_error().unwrap().is_none());
     assert!(socket2.take_error().unwrap().is_none());
@@ -435,13 +401,8 @@ fn unconnected_udp_socket_connected_methods() {
     // Receive methods don't require the socket to be connected, you just won't
     // know the sender.
     let mut buf = [0; 20];
-    let n = socket2.peek(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
-
-    let n = socket2.recv(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
+    expect_read!(socket2.peek(&mut buf), DATA1);
+    expect_read!(socket2.recv(&mut buf), DATA1);
 
     assert!(socket1.take_error().unwrap().is_none());
     assert!(socket2.take_error().unwrap().is_none());
@@ -498,15 +459,8 @@ fn connected_udp_socket_unconnected_methods() {
     );
 
     let mut buf = [0; 20];
-    let (n, got_address1) = socket3.peek_from(&mut buf).unwrap();
-    assert_eq!(n, DATA2.len());
-    assert_eq!(buf[..n], DATA2[..]);
-    assert_eq!(got_address1, address2);
-
-    let (n, got_address2) = socket3.recv_from(&mut buf).unwrap();
-    assert_eq!(n, DATA2.len());
-    assert_eq!(buf[..n], DATA2[..]);
-    assert_eq!(got_address2, address2);
+    expect_read!(socket3.peek_from(&mut buf), DATA2, address2);
+    expect_read!(socket3.recv_from(&mut buf), DATA2, address2);
 
     assert!(socket1.take_error().unwrap().is_none());
     assert!(socket2.take_error().unwrap().is_none());
@@ -575,9 +529,7 @@ fn udp_socket_reregister() {
     );
 
     let mut buf = [0; 20];
-    let (n, _) = socket.recv_from(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
+    expect_read!(socket.recv_from(&mut buf), DATA1, __anywhere);
 
     thread_handle.join().expect("unable to join thread");
 }
@@ -605,9 +557,7 @@ fn udp_socket_no_events_after_deregister() {
 
     // But we do expect a packet to be send.
     let mut buf = [0; 20];
-    let (n, _) = socket.recv_from(&mut buf).unwrap();
-    assert_eq!(n, DATA1.len());
-    assert_eq!(buf[..n], DATA1[..]);
+    expect_read!(socket.recv_from(&mut buf), DATA1, __anywhere);
 
     thread_handle.join().expect("unable to join thread");
 }

diff --git a/tests/unix_datagram.rs b/tests/unix_datagram.rs
@@ -16,8 +16,6 @@ use util::{
 
 const DATA1: &[u8] = b"Hello same host!";
 const DATA2: &[u8] = b"Why hello mio!";
-const DATA1_LEN: usize = DATA1.len();
-const DATA2_LEN: usize = DATA2.len();
 const DEFAULT_BUF_SIZE: usize = 64;
 const TEST_DIR: &str = "mio_unix_datagram_tests";
 const TOKEN_1: Token = Token(0);
@@ -162,14 +160,8 @@ fn unix_datagram_pair() {
         ],
     );
 
-    let read = datagram2.recv(&mut buf).unwrap();
-    assert_would_block(datagram2.recv(&mut buf));
-    assert_eq!(read, DATA1_LEN);
-    assert_eq!(&buf[..read], DATA1);
-
-    let read = datagram1.recv(&mut buf).unwrap();
-    assert_eq!(read, DATA2_LEN);
-    assert_eq!(&buf[..read], DATA2);
+    expect_read!(datagram1.recv(&mut buf), DATA2);
+    expect_read!(datagram2.recv(&mut buf), DATA1);
 
     assert!(datagram1.take_error().unwrap().is_none());
     assert!(datagram2.take_error().unwrap().is_none());
@@ -223,13 +215,7 @@ fn unix_datagram_try_clone() {
         ],
     );
 
-    let (read, from_addr1) = datagram1.recv_from(&mut buf).unwrap();
-    assert_eq!(read, DATA1_LEN);
-    assert_eq!(buf[..read], DATA1[..]);
-    assert_eq!(
-        from_addr1.as_pathname().expect("failed to get pathname"),
-        path2
-    );
+    expect_read!(datagram1.recv_from(&mut buf), DATA1, path: path2);
     assert_would_block(datagram2.recv_from(&mut buf));
 
     assert!(datagram1.take_error().unwrap().is_none());
@@ -397,21 +383,8 @@ fn smoke_test_unconnected(datagram1: UnixDatagram, datagram2: UnixDatagram) {
         ],
     );
 
-    let (read, from_addr1) = datagram1.recv_from(&mut buf).unwrap();
-    assert_eq!(read, DATA2_LEN);
-    assert_eq!(buf[..read], DATA2[..]);
-    assert_eq!(
-        from_addr1.as_pathname().expect("failed to get pathname"),
-        path2
-    );
-
-    let (read, from_addr2) = datagram2.recv_from(&mut buf).unwrap();
-    assert_eq!(read, DATA1_LEN);
-    assert_eq!(buf[..read], DATA1[..]);
-    assert_eq!(
-        from_addr2.as_pathname().expect("failed to get pathname"),
-        path1
-    );
+    expect_read!(datagram1.recv_from(&mut buf), DATA2, path: path2);
+    expect_read!(datagram2.recv_from(&mut buf), DATA1, path: path1);
 
     assert!(datagram1.take_error().unwrap().is_none());
     assert!(datagram2.take_error().unwrap().is_none());
@@ -471,13 +444,8 @@ fn smoke_test_connected(datagram1: UnixDatagram, datagram2: UnixDatagram) {
         ],
     );
 
-    let read = datagram1.recv(&mut buf).unwrap();
-    assert_eq!(read, DATA2_LEN);
-    assert_eq!(buf[..read], DATA2[..]);
-
-    let read = datagram2.recv(&mut buf).unwrap();
-    assert_eq!(read, DATA1_LEN);
-    assert_eq!(buf[..read], DATA1[..]);
+    expect_read!(datagram1.recv(&mut buf), DATA2);
+    expect_read!(datagram2.recv(&mut buf), DATA1);
 
     assert!(datagram1.take_error().unwrap().is_none());
     assert!(datagram2.take_error().unwrap().is_none());