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 12, 2019 · dbbdca9 · dbbdca9
1 parent 67bdad5
commit dbbdca9
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 108 deletions.
diff --git a/tests/tcp.rs b/tests/tcp.rs
@@ -12,6 +12,7 @@ use net2::TcpStreamExt;
 use mio::net::{TcpListener, TcpStream};
 use mio::{Events, Interests, Poll, Token};
 
+#[macro_use]
 mod util;
 
 use util::{assert_send, assert_sync, expect_no_events, init, init_with_poll};
@@ -811,8 +812,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]
@@ -867,8 +867,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]);
 }
 
 #[test]
@@ -906,6 +905,5 @@ 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]);
 }
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/uds.rs b/tests/uds.rs
@@ -82,9 +82,7 @@ where
         vec![ExpectEvent::new(LOCAL, Interests::READABLE)],
     );
 
-    let read = local.read(&mut buf).unwrap();
-    assert_eq!(read, DATA1_LEN);
-    assert_eq!(&buf[..read], DATA1);
+    expect_read!(local.read(&mut buf), DATA1);
 
     assert!(local.take_error().unwrap().is_none());
 
@@ -282,9 +280,7 @@ fn try_clone() {
         vec![ExpectEvent::new(LOCAL_CLONE, Interests::READABLE)],
     );
 
-    let read = local_2.read(&mut buf).unwrap();
-    assert_eq!(read, DATA1_LEN);
-    assert_eq!(&buf[..read], DATA1);
+    expect_read!(local_2.read(&mut buf), DATA1);
 
     // Close the connection to allow the remote to shutdown
     drop(local_2);
@@ -339,8 +335,7 @@ fn shutdown_read() {
     ))]
     {
         let mut buf = [0; DEFAULT_BUF_SIZE];
-        let read = local.read(&mut buf).unwrap();
-        assert_eq!(read, 0);
+        expect_read!(local.read(&mut buf), &[]);
     }
 
     // Close the connection to allow the remote to shutdown
@@ -398,9 +393,7 @@ fn shutdown_write() {
 
     // Read should be ok
     let mut buf = [0; DEFAULT_BUF_SIZE];
-    let read = local.read(&mut buf).unwrap();
-    assert_eq!(read, DATA1_LEN);
-    assert_eq!(&buf[..read], DATA1);
+    expect_read!(local.read(&mut buf), DATA1);
 
     // Close the connection to allow the remote to shutdown
     drop(local);
@@ -455,8 +448,7 @@ fn shutdown_both() {
     ))]
     {
         let mut buf = [0; DEFAULT_BUF_SIZE];
-        let read = local.read(&mut buf).unwrap();
-        assert_eq!(read, 0);
+        expect_read!(local.read(&mut buf), &[]);
     }
 
     let err = local.write(DATA2).unwrap_err();

diff --git a/tests/util/mod.rs b/tests/util/mod.rs
@@ -203,10 +203,45 @@ pub fn any_local_ipv6_address() -> SocketAddr {
 /// Usage: `checked_write!(stream.write(&DATA));`
 /// Also works for send(_to): `checked_write!(socket.send_to(DATA, address))`.
 macro_rules! checked_write {
-    ($socket: ident . $method: ident ( $data: expr $(, $arg: expr)* ) ) => {
+    ($socket: ident . $method: ident ( $data: expr $(, $arg: expr)* ) ) => {{
         let data = $data;
         let n = $socket.$method($data $(, $arg)*)
             .expect("unable to write to socket");
         assert_eq!(n, data.len(), "short write");
-    };
+    }};
+}
+
+/// A checked {write, send, send_to} macro that ensures the entire buffer is
+/// written.
+///
+/// Usage: `expect_read!(stream.read(&mut buf), DATA);` writes into `buf`,
+/// comparing to `DATA`.
+/// Also works for recv(_from): `expect_read!(socket.recv_from(&mut buf), DATA, address)`.
+macro_rules! expect_read {
+    ($socket: ident . $method: ident ( $buf: expr $(, $arg: expr)* ), $expected: expr) => {{
+        let n = $socket.$method($buf $(, $arg)*)
+            .expect("unable to read from socket");
+        let expected = $expected;
+        assert_eq!(n, expected.len());
+        assert_eq!(&$buf[..n], expected);
+    }};
+    // TODO: change the call sites to check the source address.
+    // Support for recv_from and peek_from, without checking the address.
+    ($socket: ident . $method: ident ( $buf: expr $(, $arg: expr)* ), $expected: expr, __anywhere) => {{
+        let (n, _address) = $socket.$method($buf $(, $arg)*)
+            .expect("unable to read from socket");
+        let expected = $expected;
+        assert_eq!(n, expected.len());
+        assert_eq!(&$buf[..n], expected);
+    }};
+    // Support for recv_from and peek_from.
+    ($socket: ident . $method: ident ( $buf: expr $(, $arg: expr)* ), $expected: expr, $source: expr) => {{
+        let (n, address) = $socket.$method($buf $(, $arg)*)
+            .expect("unable to read from socket");
+        let expected = $expected;
+        let source = $source;
+        assert_eq!(n, expected.len());
+        assert_eq!(&$buf[..n], expected);
+        assert_eq!(address, source);
+    }};
 }