unify get_value_from_channel/expect_recv

Author: ImplOfAnImpl

dns_server/src/crawler_p2p/crawler/tests/mod.rs

@@ -570,6 +570,7 @@ fn dont_connect_to_initially_banned_peer(#[case] seed: Seed) {
     crawler.assert_banned_addresses(&[(node1.as_bannable(), ban_end_time)]);
 }
 
+// Check that a peer is banned on CrawlerEvent::ConnectionError.
 #[rstest]
 #[trace]
 #[case(Seed::from_entropy())]

dns_server/src/crawler_p2p/crawler_manager/tests/mod.rs

@@ -22,7 +22,7 @@ use p2p::{
     config::{BanDuration, BanThreshold},
     types::socket_address::SocketAddress,
 };
-use p2p_test_utils::{assert_no_value_in_channel, get_value_from_channel};
+use p2p_test_utils::{expect_no_recv, expect_recv};
 
 use crate::{
     crawler_p2p::crawler_manager::tests::mock_manager::{
@@ -41,15 +41,15 @@ async fn basic() {
     state.node_online(node1);
     advance_time(&mut crawler, &time_getter, Duration::from_secs(60), 60).await;
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node1.socket_addr().ip())
     );
 
     // Node goes offline, DNS record removed
     state.node_offline(node1);
     advance_time(&mut crawler, &time_getter, Duration::from_secs(60), 60).await;
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::DelAddress(node1.socket_addr().ip())
     );
 }
@@ -72,7 +72,7 @@ async fn long_offline() {
     state.node_online(node1);
     advance_time(&mut crawler, &time_getter, Duration::from_secs(60), 24 * 60).await;
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node1.socket_addr().ip())
     );
 }
@@ -90,21 +90,21 @@ async fn announced_online() {
 
     advance_time(&mut crawler, &time_getter, Duration::from_secs(60), 60).await;
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node1.socket_addr().ip())
     );
 
     state.announce_address(node1, node2);
     advance_time(&mut crawler, &time_getter, Duration::from_secs(60), 60).await;
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node2.socket_addr().ip())
     );
 
     state.announce_address(node2, node3);
     advance_time(&mut crawler, &time_getter, Duration::from_secs(60), 60).await;
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node3.socket_addr().ip())
     );
 
@@ -121,7 +121,7 @@ async fn announced_offline() {
 
     advance_time(&mut crawler, &time_getter, Duration::from_secs(60), 60).await;
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node1.socket_addr().ip())
     );
     assert_eq!(state.connection_attempts.lock().unwrap().len(), 1);
@@ -136,7 +136,7 @@ async fn announced_offline() {
     state.announce_address(node1, node2);
     advance_time(&mut crawler, &time_getter, Duration::from_secs(60), 24 * 60).await;
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node2.socket_addr().ip())
     );
     assert_eq!(state.connection_attempts.lock().unwrap().len(), 3);
@@ -164,14 +164,14 @@ async fn private_ip() {
 
     // Check that only nodes with public addresses and on the default port are added to DNS
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node1.socket_addr().ip())
     );
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node2.socket_addr().ip())
     );
-    assert_no_value_in_channel(&mut command_rx).await;
+    expect_no_recv!(command_rx);
 
     // Check that all reachable nodes are stored in the DB
     assert_known_addresses(&crawler, &[node1, node2, node3, node4, node5, node6]);
@@ -202,14 +202,14 @@ async fn ban_unban() {
 
     // Only normal nodes are added to DNS
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node1.socket_addr().ip())
     );
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node3.socket_addr().ip())
     );
-    assert_no_value_in_channel(&mut command_rx).await;
+    expect_no_recv!(command_rx);
 
     // node2 is banned
     assert_banned_addresses(&crawler, &[(node2.as_bannable(), node2_ban_end_time)]);
@@ -226,7 +226,7 @@ async fn ban_unban() {
 
     // Check that it's been removed from DNS.
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::DelAddress(node1.socket_addr().ip())
     );
 
@@ -252,7 +252,7 @@ async fn ban_unban() {
             (node2.as_bannable(), node2_ban_end_time),
         ],
     );
-    assert_no_value_in_channel(&mut command_rx).await;
+    expect_no_recv!(command_rx);
 
     // Wait enough time for node2 to be unbanned.
     let time_until_node2_unban =
@@ -262,7 +262,7 @@ async fn ban_unban() {
     // node2 is no longer banned; its address has been added to DNS.
     assert_banned_addresses(&crawler, &[(node1.as_bannable(), node1_ban_end_time)]);
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node2.socket_addr().ip())
     );
 
@@ -274,7 +274,7 @@ async fn ban_unban() {
     // node1 is no longer banned; its address has been added to DNS.
     assert_banned_addresses(&crawler, &[]);
     assert_eq!(
-        get_value_from_channel(&mut command_rx).await.unwrap(),
+        expect_recv!(command_rx),
         DnsServerCommand::AddAddress(node1.socket_addr().ip())
     );
 }

p2p/p2p-test-utils/src/lib.rs

@@ -15,7 +15,7 @@
 
 #![allow(clippy::unwrap_used)]
 
-use std::{fmt::Debug, sync::Arc, time::Duration};
+use std::{sync::Arc, time::Duration};
 
 use chainstate::{
     chainstate_interface::ChainstateInterface, make_chainstate, ChainstateConfig,
@@ -30,7 +30,6 @@ use common::{
 use mempool::{MempoolHandle, MempoolSubsystemInterface};
 use subsystem::manager::{ManagerJoinHandle, ShutdownTrigger};
 use test_utils::mock_time_getter::mocked_time_getter_milliseconds;
-use tokio::{sync::mpsc::UnboundedReceiver, time};
 use utils::atomics::SeqCstAtomicU64;
 
 pub fn start_subsystems(
@@ -122,12 +121,40 @@ pub const LONG_TIMEOUT: Duration = Duration::from_secs(60);
 /// A short timeout for events that shouldn't occur.
 pub const SHORT_TIMEOUT: Duration = Duration::from_millis(500);
 
-pub async fn get_value_from_channel<T>(rx: &mut UnboundedReceiver<T>) -> Option<T> {
-    time::timeout(LONG_TIMEOUT, rx.recv())
-        .await
-        .expect("Failed to receive value in time")
+/// Await for the specified future for some reasonably big amount of time; panic if the timeout
+/// is reached.
+// Note: this is implemented as a macro until #[track_caller] works correctly with async functions
+// (needed to print the caller location if 'unwrap' fails). Same for the other macros below.
+#[macro_export]
+macro_rules! expect_future_val {
+    ($fut:expr) => {
+        tokio::time::timeout($crate::LONG_TIMEOUT, $fut)
+            .await
+            .expect("Failed to receive value in time")
+    };
 }
 
-pub async fn assert_no_value_in_channel<T: Debug>(rx: &mut UnboundedReceiver<T>) {
-    time::timeout(SHORT_TIMEOUT, rx.recv()).await.unwrap_err();
+/// Await for the specified future for a short time, expecting a timeout.
+#[macro_export]
+macro_rules! expect_no_future_val {
+    ($fut:expr) => {
+        tokio::time::timeout($crate::SHORT_TIMEOUT, $fut).await.unwrap_err();
+    };
+}
+
+/// Try receiving a message from the tokio channel; panic if the channel is closed or the timeout
+/// is reached.
+#[macro_export]
+macro_rules! expect_recv {
+    ($rx:expr) => {
+        $crate::expect_future_val!($rx.recv()).unwrap()
+    };
+}
+
+/// Try receiving a message from the tokio channel; expect that a timeout is reached.
+#[macro_export]
+macro_rules! expect_no_recv {
+    ($rx:expr) => {
+        $crate::expect_no_future_val!($rx.recv())
+    };
 }

p2p/src/peer_manager/tests/mod.rs

@@ -22,6 +22,7 @@ mod utils;
 
 use std::{sync::Arc, time::Duration};
 
+use p2p_test_utils::expect_recv;
 use tokio::sync::{mpsc, oneshot};
 
 use ::utils::atomics::SeqCstAtomicBool;
@@ -35,7 +36,6 @@ use tokio::{
 };
 
 use crate::{
-    expect_recv,
     interface::types::ConnectedPeer,
     message::{PeerManagerMessage, PingRequest, PingResponse},
     net::{

p2p/src/peer_manager/tests/ping.rs

@@ -16,12 +16,11 @@
 use std::{sync::Arc, time::Duration};
 
 use common::{chain::config, primitives::user_agent::mintlayer_core_user_agent};
-use p2p_test_utils::P2pBasicTestTimeGetter;
+use p2p_test_utils::{expect_recv, P2pBasicTestTimeGetter};
 use test_utils::{assert_matches, assert_matches_return_val};
 
 use crate::{
     config::{NodeType, P2pConfig},
-    expect_recv,
     message::{PeerManagerMessage, PingRequest, PingResponse},
     net::{
         default_backend::{

p2p/src/sync/tests/helpers/mod.rs

@@ -17,7 +17,7 @@ use std::{collections::BTreeMap, panic, sync::Arc};
 
 use async_trait::async_trait;
 use crypto::random::Rng;
-use p2p_test_utils::{assert_no_value_in_channel, get_value_from_channel, SHORT_TIMEOUT};
+use p2p_test_utils::{expect_future_val, expect_no_recv, expect_recv, SHORT_TIMEOUT};
 use p2p_types::socket_address::SocketAddress;
 use test_utils::random::Seed;
 use tokio::{
@@ -212,7 +212,7 @@ impl TestNode {
 
     /// Receives a message from the sync manager.
     pub async fn message(&mut self) -> (PeerId, SyncMessage) {
-        get_value_from_channel(&mut self.sync_msg_receiver).await.unwrap()
+        expect_recv!(&mut self.sync_msg_receiver)
     }
 
     /// Try to receive a message from the sync manager.
@@ -231,7 +231,7 @@ impl TestNode {
 
     /// Panics if the sync manager returns an error.
     pub async fn assert_no_error(&mut self) {
-        assert_no_value_in_channel(&mut self.error_receiver).await;
+        expect_no_recv!(self.error_receiver);
     }
 
     /// Receives the `AdjustPeerScore` event from the peer manager.
@@ -500,9 +500,7 @@ struct SyncingEventReceiverMock {
 #[async_trait]
 impl SyncingEventReceiver for SyncingEventReceiverMock {
     async fn poll_next(&mut self) -> Result<SyncingEvent> {
-        get_value_from_channel(&mut self.events_receiver)
-            .await
-            .ok_or(P2pError::ChannelClosed)
+        expect_future_val!(self.events_receiver.recv()).ok_or(P2pError::ChannelClosed)
     }
 }
 

p2p/src/testing_utils.rs

@@ -225,17 +225,6 @@ pub fn peerdb_inmemory_store() -> PeerDbStorageImpl<storage::inmemory::InMemory>
     PeerDbStorageImpl::new(storage).unwrap()
 }
 
-/// Receive a message from the tokio channel.
-/// Panics if the channel is closed or no message received in 10 seconds.
-#[macro_export]
-macro_rules! expect_recv {
-    // Implemented as a macro until #[track_caller] works correctly with async functions
-    // (needed to print the caller location if unwraps fail)
-    ($x:expr) => {
-        tokio::time::timeout(Duration::from_secs(10), $x.recv()).await.unwrap().unwrap()
-    };
-}
-
 pub fn test_p2p_config() -> P2pConfig {
     P2pConfig {
         bind_addresses: Default::default(),