Repository: mintlayer/mintlayer-core
Pull request #655 (2023-01-27 13:02:57)

Use TimeGetter in p2p


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 3e7d45918..aa1a1e156 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -600,7 +600,6 @@ dependencies = [
  "generic-array",
  "hex",
  "itertools",
- "lazy_static",
  "logging",
  "merkletree",
  "parity-scale-codec",
diff --git a/common/Cargo.toml b/common/Cargo.toml
index f69cf35ac..f192b38f9 100644
--- a/common/Cargo.toml
+++ b/common/Cargo.toml
@@ -17,7 +17,6 @@ typename = { path = "../utils/typename" }
 anyhow = "1.0"
 bech32 = "0.9"
 generic-array = "0.14"
-lazy_static = "1.4"
 merkletree = "0.22"
 parity-scale-codec = "3.1"
 serde = "1.0"
diff --git a/common/src/primitives/time.rs b/common/src/primitives/time.rs
index 86f3b5250..377aa0b3c 100644
--- a/common/src/primitives/time.rs
+++ b/common/src/primitives/time.rs
@@ -13,9 +13,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#![allow(unused, dead_code)]
-use lazy_static::lazy_static;
-use logging::log;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, SystemTime};
 
@@ -28,9 +25,7 @@ fn duration_from_int(v: u64) -> Duration {
     Duration::from_millis(v)
 }
 
-lazy_static! {
-    static ref TIME_SOURCE: AtomicU64 = Default::default();
-}
+static TIME_SOURCE: AtomicU64 = AtomicU64::new(0);
 
 /// Either gets the current time or panics
 pub fn get() -> Duration {
@@ -67,13 +62,15 @@ pub fn set(now: Duration) -> Result<(), Box<dyn std::error::Error>> {
 
 #[cfg(test)]
 mod tests {
+    use logging::log;
+
     use super::*;
 
     #[test]
     #[serial_test::serial]
     fn test_time() {
         logging::init_logging::<&std::path::Path>(None);
-        set(Duration::from_secs(1337));
+        set(Duration::from_secs(1337)).unwrap();
 
         log::info!("p2p time: {}", get().as_secs());
         std::thread::sleep(Duration::from_secs(1));
diff --git a/node/src/config_files/p2p.rs b/node/src/config_files/p2p.rs
index 7d7050e73..a8c6b0ce4 100644
--- a/node/src/config_files/p2p.rs
+++ b/node/src/config_files/p2p.rs
@@ -78,8 +78,6 @@ impl From<P2pConfigFile> for P2pConfig {
                 .into(),
             node_type: c.node_type.map(Into::into).into(),
             allow_discover_private_ips: Default::default(),
-            heartbeat_interval_min: Default::default(),
-            heartbeat_interval_max: Default::default(),
         }
     }
 }
diff --git a/node/src/runner.rs b/node/src/runner.rs
index 25518fd73..bff6be8e9 100644
--- a/node/src/runner.rs
+++ b/node/src/runner.rs
@@ -71,7 +71,7 @@ pub async fn initialize(
     });
 
     // P2P subsystem
-    // TODO: Replace Lmdb with Sqlite backend when it's read
+    // TODO: Replace Lmdb with Sqlite backend when it's ready
     let peerdb_storage = PeerDbStorageImpl::new(storage_lmdb::Lmdb::new(
         node_config.datadir.join("peerdb-lmdb"),
     ))?;
@@ -82,6 +82,7 @@ pub async fn initialize(
             Arc::new(node_config.p2p.into()),
             chainstate.clone(),
             mempool.clone(),
+            Default::default(),
             peerdb_storage,
         )
         .await
diff --git a/p2p/Cargo.toml b/p2p/Cargo.toml
index 3b8924a5e..7581b7afa 100644
--- a/p2p/Cargo.toml
+++ b/p2p/Cargo.toml
@@ -44,6 +44,7 @@ p2p-test-utils = { path = "p2p-test-utils" }
 p2p-backend-test-suite = { path = "backend-test-suite" }
 
 portpicker = "0.1"
+tokio = { version = "1", features = ["test-util"] }
 
 [[test]]
 name = "backend_tcp"
diff --git a/p2p/src/config.rs b/p2p/src/config.rs
index 44e1afd43..a405c9934 100644
--- a/p2p/src/config.rs
+++ b/p2p/src/config.rs
@@ -31,8 +31,6 @@ make_config_setting!(
 );
 make_config_setting!(NodeTypeSetting, NodeType, NodeType::Full);
 make_config_setting!(AllowDiscoverPrivateIps, bool, false);
-make_config_setting!(HeartbeatIntervalMin, Duration, Duration::from_secs(5));
-make_config_setting!(HeartbeatIntervalMax, Duration, Duration::from_secs(30));
 
 /// A node type.
 #[derive(Debug, Copy, Clone)]
@@ -76,8 +74,4 @@ pub struct P2pConfig {
     pub node_type: NodeTypeSetting,
     /// Allow announcing and discovering local and private IPs. Should be used for testing only.
     pub allow_discover_private_ips: AllowDiscoverPrivateIps,
-    /// Lower bound for how often `PeerManager::heartbeat` is called. Should be used for testing only.
-    pub heartbeat_interval_min: HeartbeatIntervalMin,
-    /// Upper bound for how often `PeerManager::heartbeat` is called. Should be used for testing only.
-    pub heartbeat_interval_max: HeartbeatIntervalMin,
 }
diff --git a/p2p/src/lib.rs b/p2p/src/lib.rs
index d43a9afd4..26887ff80 100644
--- a/p2p/src/lib.rs
+++ b/p2p/src/lib.rs
@@ -35,7 +35,7 @@ use tap::TapFallible;
 use tokio::sync::mpsc;
 
 use chainstate::chainstate_interface;
-use common::chain::ChainConfig;
+use common::{chain::ChainConfig, time_getter::TimeGetter};
 use logging::log;
 
 use crate::{
@@ -78,6 +78,7 @@ where
         p2p_config: Arc<P2pConfig>,
         chainstate_handle: subsystem::Handle<Box<dyn chainstate_interface::ChainstateInterface>>,
         _mempool_handle: mempool::MempoolHandle,
+        time_getter: TimeGetter,
         peerdb_storage: S,
     ) -> crate::Result<Self> {
         let bind_addresses = p2p_config
@@ -116,6 +117,7 @@ where
             conn,
             rx_peer_manager,
             tx_p2p_sync,
+            time_getter,
             peerdb_storage,
         )?;
         tokio::spawn(async move {
@@ -158,6 +160,7 @@ pub async fn make_p2p<S: PeerDbStorage + 'static>(
     p2p_config: Arc<P2pConfig>,
     chainstate_handle: subsystem::Handle<Box<dyn chainstate_interface::ChainstateInterface>>,
     mempool_handle: mempool::MempoolHandle,
+    time_getter: TimeGetter,
     peerdb_storage: S,
 ) -> Result<Box<dyn P2pInterface>> {
     let stream_adapter = NoiseEncryptionAdapter::gen_new();
@@ -170,6 +173,7 @@ pub async fn make_p2p<S: PeerDbStorage + 'static>(
         p2p_config,
         chainstate_handle,
         mempool_handle,
+        time_getter,
         peerdb_storage,
     )
     .await?;
diff --git a/p2p/src/peer_manager/mod.rs b/p2p/src/peer_manager/mod.rs
index 10ace4e91..f1de22f32 100644
--- a/p2p/src/peer_manager/mod.rs
+++ b/p2p/src/peer_manager/mod.rs
@@ -27,13 +27,16 @@ pub mod peerdb;
 use std::{
     collections::{HashMap, HashSet},
     sync::Arc,
-    time::Instant,
+    time::Duration,
 };
 
-use tokio::sync::{mpsc, oneshot};
+use tokio::{
+    sync::{mpsc, oneshot},
+    time::Instant,
+};
 
 use chainstate::ban_score::BanScore;
-use common::{chain::ChainConfig, primitives::semver::SemVer};
+use common::{chain::ChainConfig, primitives::semver::SemVer, time_getter::TimeGetter};
 use logging::log;
 use utils::ensure;
 
@@ -60,6 +63,11 @@ use self::{global_ip::IsGlobalIp, peerdb::storage::PeerDbStorage};
 /// Maximum number of connections the [`PeerManager`] is allowed to have open
 const MAX_ACTIVE_CONNECTIONS: usize = 128;
 
+/// Lower bound for how often [`PeerManager::heartbeat()`] is called
+const PEER_MGR_HEARTBEAT_INTERVAL_MIN: Duration = Duration::from_secs(5);
+/// Upper bound for how often [`PeerManager::heartbeat()`] is called
+const PEER_MGR_HEARTBEAT_INTERVAL_MAX: Duration = Duration::from_secs(30);
+
 /// How many addresses are allowed to be sent
 const MAX_ADDRESS_COUNT: usize = 1000;
 
@@ -112,9 +120,10 @@ where
         handle: T::ConnectivityHandle,
         rx_peer_manager: mpsc::UnboundedReceiver<PeerManagerEvent<T>>,
         tx_sync: mpsc::UnboundedSender<SyncControlEvent<T>>,
+        time_getter: TimeGetter,
         peerdb_storage: S,
     ) -> crate::Result<Self> {
-        let peerdb = peerdb::PeerDb::new(Arc::clone(&p2p_config), peerdb_storage)?;
+        let peerdb = peerdb::PeerDb::new(Arc::clone(&p2p_config), time_getter, peerdb_storage)?;
         Ok(Self {
             peer_connectivity_handle: handle,
             rx_peer_manager,
@@ -123,7 +132,7 @@ where
             pending: HashMap::new(),
             chain_config,
             p2p_config,
-            last_heartbeat: Instant::now(),
+            last_heartbeat: tokio::time::Instant::now(),
             announced_addresses: HashMap::new(),
         })
     }
@@ -704,12 +713,12 @@ where
                 event_res = self.peer_connectivity_handle.poll_next() => {
                     self.handle_connectivity_event_result(event_res).await?;
                 },
-                _event = tokio::time::sleep(*self.p2p_config.heartbeat_interval_max) => {}
+                _event = tokio::time::sleep(PEER_MGR_HEARTBEAT_INTERVAL_MAX) => {}
             }
 
             // finally update peer manager state
-            let now = Instant::now();
-            if now.duration_since(self.last_heartbeat) > *self.p2p_config.heartbeat_interval_min {
+            let now = tokio::time::Instant::now();
+            if now.duration_since(self.last_heartbeat) > PEER_MGR_HEARTBEAT_INTERVAL_MIN {
                 self.heartbeat().await?;
                 self.last_heartbeat = now;
             }
diff --git a/p2p/src/peer_manager/peerdb/mod.rs b/p2p/src/peer_manager/peerdb/mod.rs
index 36a483219..d2c3aaa16 100644
--- a/p2p/src/peer_manager/peerdb/mod.rs
+++ b/p2p/src/peer_manager/peerdb/mod.rs
@@ -31,9 +31,10 @@ pub mod storage_impl;
 use std::{
     collections::{BTreeMap, BTreeSet},
     sync::Arc,
-    time::{Duration, SystemTime, UNIX_EPOCH},
+    time::Duration,
 };
 
+use common::time_getter::TimeGetter;
 use crypto::random::{make_pseudo_rng, SliceRandom};
 use logging::log;
 
@@ -97,11 +98,17 @@ pub struct PeerDb<T: NetworkingService, S> {
     /// when `current_time > ban_duration`.
     banned_addresses: BTreeMap<T::BannableAddress, Duration>,
 
+    time_getter: TimeGetter,
+
     storage: S,
 }
 
 impl<T: NetworkingService, S: PeerDbStorage> PeerDb<T, S> {
-    pub fn new(p2p_config: Arc<config::P2pConfig>, storage: S) -> crate::Result<Self> {
+    pub fn new(
+        p2p_config: Arc<config::P2pConfig>,
+        time_getter: TimeGetter,
+        storage: S,
+    ) -> crate::Result<Self> {
         let added_nodes = p2p_config
             .added_nodes
             .iter()
@@ -137,6 +144,7 @@ impl<T: NetworkingService, S: PeerDbStorage> PeerDb<T, S> {
             known_addresses,
             banned_addresses,
             p2p_config,
+            time_getter,
             storage,
         })
     }
@@ -190,10 +198,7 @@ impl<T: NetworkingService, S: PeerDbStorage> PeerDb<T, S> {
     pub fn is_address_banned(&mut self, address: &T::BannableAddress) -> crate::Result<bool> {
         if let Some(banned_till) = self.banned_addresses.get(address) {
             // Check if the ban has expired.
-            let now = SystemTime::now()
-                .duration_since(UNIX_EPOCH)
-                // This can fail only if `SystemTime::now()` returns the time before `UNIX_EPOCH`.
-                .expect("Invalid system time");
+            let now = self.time_getter.get_time();
             if now > *banned_till {
                 self.banned_addresses.remove(address);
                 let mut tx = self.storage.transaction_rw()?;
@@ -288,11 +293,7 @@ impl<T: NetworkingService, S: PeerDbStorage> PeerDb<T, S> {
     fn ban_peer(&mut self, peer_id: &T::PeerId) -> crate::Result<()> {
         if let Some(peer) = self.peers.remove(peer_id) {
             let bannable_address = peer.address.as_bannable();
-            let ban_till = SystemTime::now()
-                .duration_since(UNIX_EPOCH)
-                // This can fail only if `SystemTime::now()` returns the time before `UNIX_EPOCH`.
-                .expect("Invalid system time")
-                + *self.p2p_config.ban_duration;
+            let ban_till = self.time_getter.get_time() + *self.p2p_config.ban_duration;
             let mut tx = self.storage.transaction_rw()?;
             tx.add_banned_address(&bannable_address.to_string(), ban_till)?;
             tx.commit()?;


Test patch:
diff --git a/p2p/backend-test-suite/src/block_announcement.rs b/p2p/backend-test-suite/src/block_announcement.rs
index 3338453bb..621d5fb1c 100644
--- a/p2p/backend-test-suite/src/block_announcement.rs
+++ b/p2p/backend-test-suite/src/block_announcement.rs
@@ -137,8 +137,6 @@ where
         outbound_connection_timeout: Default::default(),
         node_type: NodeType::Inactive.into(),
         allow_discover_private_ips: Default::default(),
-        heartbeat_interval_min: Default::default(),
-        heartbeat_interval_max: Default::default(),
     });
     let (mut conn1, mut sync1) = N::start(
         T::make_transport(),
diff --git a/p2p/backend-test-suite/src/peer_manager_peerdb.rs b/p2p/backend-test-suite/src/peer_manager_peerdb.rs
index 7bf9cff5a..dd969df4b 100644
--- a/p2p/backend-test-suite/src/peer_manager_peerdb.rs
+++ b/p2p/backend-test-suite/src/peer_manager_peerdb.rs
@@ -44,8 +44,12 @@ where
     N: NetworkingService<PeerId = PeerId>,
     A: RandomAddressMaker<Address = N::Address>,
 {
-    let mut peerdb =
-        PeerDb::<N, _>::new(Arc::new(P2pConfig::default()), peerdb_inmemory_store()).unwrap();
+    let mut peerdb = PeerDb::<N, _>::new(
+        Arc::new(P2pConfig::default()),
+        Default::default(),
+        peerdb_inmemory_store(),
+    )
+    .unwrap();
 
     let peer_id = add_active_peer::<T, N, A, _>(&mut peerdb);
     let address = peerdb.peer_address(&peer_id).unwrap().as_bannable();
@@ -66,10 +70,13 @@ where
         outbound_connection_timeout: Default::default(),
         node_type: Default::default(),
         allow_discover_private_ips: Default::default(),
-        heartbeat_interval_min: Default::default(),
-        heartbeat_interval_max: Default::default(),
     };
-    let mut peerdb = PeerDb::<N, _>::new(Arc::new(config), peerdb_inmemory_store()).unwrap();
+    let mut peerdb = PeerDb::<N, _>::new(
+        Arc::new(config),
+        Default::default(),
+        peerdb_inmemory_store(),
+    )
+    .unwrap();
 
     let peer_id = add_active_peer::<T, N, A, _>(&mut peerdb);
     assert!(!peerdb.adjust_peer_score(&peer_id, 100).unwrap());
@@ -91,10 +98,13 @@ where
         outbound_connection_timeout: Default::default(),
         node_type: Default::default(),
         allow_discover_private_ips: Default::default(),
-        heartbeat_interval_min: Default::default(),
-        heartbeat_interval_max: Default::default(),
     };
-    let mut peerdb = PeerDb::<N, _>::new(Arc::new(config), peerdb_inmemory_store()).unwrap();
+    let mut peerdb = PeerDb::<N, _>::new(
+        Arc::new(config),
+        Default::default(),
+        peerdb_inmemory_store(),
+    )
+    .unwrap();
 
     let peer_id = add_active_peer::<T, N, A, _>(&mut peerdb);
     let address = peerdb.peer_address(&peer_id).unwrap().as_bannable();
@@ -117,9 +127,8 @@ where
             outbound_connection_timeout: Default::default(),
             node_type: Default::default(),
             allow_discover_private_ips: Default::default(),
-            heartbeat_interval_min: Default::default(),
-            heartbeat_interval_max: Default::default(),
         }),
+        Default::default(),
         db_store,
     )
     .unwrap();
diff --git a/p2p/src/peer_manager/tests/connections.rs b/p2p/src/peer_manager/tests/connections.rs
index 5135f09fa..561b42e75 100644
--- a/p2p/src/peer_manager/tests/connections.rs
+++ b/p2p/src/peer_manager/tests/connections.rs
@@ -24,10 +24,10 @@ use tokio::{sync::oneshot, time::timeout};
 use crate::{
     config::P2pConfig,
     net::types::Role,
-    peer_manager::tests::run_peer_manager,
+    peer_manager::tests::{get_connected_peers, run_peer_manager},
     testing_utils::{
-        connect_services, get_connectivity_event, peerdb_inmemory_store, TestTransportChannel,
-        TestTransportMaker, TestTransportNoise, TestTransportTcp,
+        connect_services, get_connectivity_event, peerdb_inmemory_store, P2pTestTimeGetter,
+        TestTransportChannel, TestTransportMaker, TestTransportNoise, TestTransportTcp,
     },
 };
 use common::chain::config;
@@ -610,6 +610,7 @@ async fn connection_timeout_rpc_notified<T>(
         conn,
         rx,
         tx_sync,
+        Default::default(),
         peerdb_inmemory_store(),
     )
     .unwrap();
@@ -679,14 +680,13 @@ where
         outbound_connection_timeout: Default::default(),
         node_type: Default::default(),
         allow_discover_private_ips: Default::default(),
-        heartbeat_interval_min: Default::default(),
-        heartbeat_interval_max: Default::default(),
     });
     let tx1 = run_peer_manager::<T>(
         A::make_transport(),
         A::make_address(),
         Arc::clone(&chain_config),
         p2p_config_1,
+        Default::default(),
     )
     .await;
 
@@ -705,14 +705,13 @@ where
         outbound_connection_timeout: Default::default(),
         node_type: Default::default(),
         allow_discover_private_ips: Default::default(),
-        heartbeat_interval_min: Default::default(),
-        heartbeat_interval_max: Default::default(),
     });
     let tx1 = run_peer_manager::<T>(
         A::make_transport(),
         A::make_address(),
         Arc::clone(&chain_config),
         p2p_config_2,
+        Default::default(),
     )
     .await;
 
@@ -759,6 +758,8 @@ where
 {
     let chain_config = Arc::new(config::create_mainnet());
 
+    let time_getter = P2pTestTimeGetter::new();
+
     // Start the first peer manager
     let p2p_config_1 = Arc::new(P2pConfig {
         bind_addresses: Default::default(),
@@ -768,20 +769,20 @@ where
         outbound_connection_timeout: Default::default(),
         node_type: Default::default(),
         allow_discover_private_ips: true.into(),
-        heartbeat_interval_min: Duration::from_secs(1).into(),
-        heartbeat_interval_max: Duration::from_secs(2).into(),
     });
     let tx1 = run_peer_manager::<T>(
         A::make_transport(),
         A::make_address(),
         Arc::clone(&chain_config),
         p2p_config_1,
+        time_getter.get_time_getter(),
     )
     .await;
 
     // Get the first peer manager's bind address
     let (rtx, rrx) = oneshot::channel();
     tx1.send(PeerManagerEvent::GetBindAddresses(rtx)).unwrap();
+
     let bind_addresses = timeout(Duration::from_secs(1), rrx).await.unwrap().unwrap();
     assert_eq!(bind_addresses.len(), 1);
 
@@ -794,14 +795,13 @@ where
         outbound_connection_timeout: Default::default(),
         node_type: Default::default(),
         allow_discover_private_ips: true.into(),
-        heartbeat_interval_min: Duration::from_secs(1).into(),
-        heartbeat_interval_max: Duration::from_secs(2).into(),
     });
     let tx2 = run_peer_manager::<T>(
         A::make_transport(),
         A::make_address(),
         Arc::clone(&chain_config),
         p2p_config_2,
+        time_getter.get_time_getter(),
     )
     .await;
 
@@ -814,66 +814,42 @@ where
         outbound_connection_timeout: Default::default(),
         node_type: Default::default(),
         allow_discover_private_ips: true.into(),
-        heartbeat_interval_min: Duration::from_secs(1).into(),
-        heartbeat_interval_max: Duration::from_secs(2).into(),
     });
     let tx3 = run_peer_manager::<T>(
         A::make_transport(),
         A::make_address(),
         Arc::clone(&chain_config),
         p2p_config_3,
+        time_getter.get_time_getter(),
     )
     .await;
 
     let started_at = Instant::now();
 
-    // The second peer should discover the third peer
+    // All peers should discover each other
     loop {
-        tokio::time::sleep(Duration::from_millis(100)).await;
-        let (rtx, rrx) = oneshot::channel();
-        tx2.send(PeerManagerEvent::GetConnectedPeers(rtx)).unwrap();
-        let connected_peers = timeout(Duration::from_secs(1), rrx).await.unwrap().unwrap();
-        if connected_peers.len() == expected_count {
+        let connected_peers_1 = get_connected_peers(&tx1).await.len();
+        let connected_peers_2 = get_connected_peers(&tx2).await.len();
+        let connected_peers_3 = get_connected_peers(&tx3).await.len();
+
+        if connected_peers_1 == expected_count
+            && connected_peers_2 == expected_count
+            && connected_peers_3 == expected_count
+        {
             break;
         }
-        assert!(
-            Instant::now().duration_since(started_at) < Duration::from_secs(10),
-            "Unexpected peer count: {}, expected: {}",
-            connected_peers.len(),
-            expected_count,
-        );
-    }
 
-    // Check the third peer manager
-    loop {
-        tokio::time::sleep(Duration::from_millis(100)).await;
-        let (rtx, rrx) = oneshot::channel();
-        tx3.send(PeerManagerEvent::GetConnectedPeers(rtx)).unwrap();
-        let connected_peers = timeout(Duration::from_secs(1), rrx).await.unwrap().unwrap();
-        if connected_peers.len() == expected_count {
-            break;
-        }
-        assert!(
-            Instant::now().duration_since(started_at) < Duration::from_secs(10),
-            "Unexpected peer count: {}, expected: {}",
-            connected_peers.len(),
-            expected_count,
-        );
-    }
+        // Without this noise handshake does complete in time for some reasons
+        tokio::time::sleep(Duration::from_millis(10)).await;
+
+        time_getter.advance_time(Duration::from_millis(1000)).await;
 
-    // And the first peer manager
-    loop {
-        tokio::time::sleep(Duration::from_millis(100)).await;
-        let (rtx, rrx) = oneshot::channel();
-        tx1.send(PeerManagerEvent::GetConnectedPeers(rtx)).unwrap();
-        let connected_peers = timeout(Duration::from_secs(1), rrx).await.unwrap().unwrap();
-        if connected_peers.len() == expected_count {
-            break;
-        }
         assert!(
             Instant::now().duration_since(started_at) < Duration::from_secs(10),
-            "Unexpected peer count: {}, expected: {}",
-            connected_peers.len(),
+            "Unexpected peer counts: {}, {}, {}, expected: {}",
+            connected_peers_1,
+            connected_peers_2,
+            connected_peers_3,
             expected_count,
         );
     }
diff --git a/p2p/src/peer_manager/tests/mod.rs b/p2p/src/peer_manager/tests/mod.rs
index 279d6ff85..11095b071 100644
--- a/p2p/src/peer_manager/tests/mod.rs
+++ b/p2p/src/peer_manager/tests/mod.rs
@@ -16,12 +16,17 @@
 mod ban;
 mod connections;
 
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
-use tokio::sync::mpsc::UnboundedSender;
+use common::time_getter::TimeGetter;
+use tokio::{
+    sync::{mpsc::UnboundedSender, oneshot},
+    time::timeout,
+};
 
 use crate::{
     event::PeerManagerEvent,
+    interface::types::ConnectedPeer,
     net::{ConnectivityService, NetworkingService},
     peer_manager::PeerManager,
     testing_utils::peerdb_inmemory_store,
@@ -35,6 +40,7 @@ async fn make_peer_manager_custom<T>(
     addr: T::Address,
     chain_config: Arc<common::chain::ChainConfig>,
     p2p_config: Arc<P2pConfig>,
+    time_getter: TimeGetter,
 ) -> (
     PeerManager<T, impl PeerDbStorage>,
     UnboundedSender<PeerManagerEvent<T>>,
@@ -62,6 +68,7 @@ where
         conn,
         rx,
         tx_sync,
+        time_getter,
         peerdb_inmemory_store(),
     )
     .unwrap();
@@ -79,8 +86,14 @@ where
     T::ConnectivityHandle: ConnectivityService<T>,
 {
     let p2p_config = Arc::new(P2pConfig::default());
-    let (peer_manager, _tx) =
-        make_peer_manager_custom::<T>(transport, addr, chain_config, p2p_config).await;
+    let (peer_manager, _tx) = make_peer_manager_custom::<T>(
+        transport,
+        addr,
+        chain_config,
+        p2p_config,
+        Default::default(),
+    )
+    .await;
     peer_manager
 }
 
@@ -89,15 +102,24 @@ async fn run_peer_manager<T>(
     addr: T::Address,
     chain_config: Arc<common::chain::ChainConfig>,
     p2p_config: Arc<P2pConfig>,
+    time_getter: TimeGetter,
 ) -> UnboundedSender<PeerManagerEvent<T>>
 where
     T: NetworkingService + 'static,
     T::ConnectivityHandle: ConnectivityService<T>,
 {
     let (mut peer_manager, tx) =
-        make_peer_manager_custom::<T>(transport, addr, chain_config, p2p_config).await;
+        make_peer_manager_custom::<T>(transport, addr, chain_config, p2p_config, time_getter).await;
     tokio::spawn(async move {
         peer_manager.run().await.unwrap();
     });
     tx
 }
+
+async fn get_connected_peers<T: NetworkingService + std::fmt::Debug>(
+    tx: &UnboundedSender<PeerManagerEvent<T>>,
+) -> Vec<ConnectedPeer> {
+    let (rtx, rrx) = oneshot::channel();
+    tx.send(PeerManagerEvent::GetConnectedPeers(rtx)).unwrap();
+    timeout(Duration::from_secs(1), rrx).await.unwrap().unwrap()
+}
diff --git a/p2p/src/sync/tests/mod.rs b/p2p/src/sync/tests/mod.rs
index 7af607a42..76890f71d 100644
--- a/p2p/src/sync/tests/mod.rs
+++ b/p2p/src/sync/tests/mod.rs
@@ -77,8 +77,6 @@ where
         outbound_connection_timeout: Default::default(),
         node_type: NodeType::Full.into(),
         allow_discover_private_ips: Default::default(),
-        heartbeat_interval_min: Default::default(),
-        heartbeat_interval_max: Default::default(),
     });
     let (conn, sync) = T::start(
         transport,
diff --git a/p2p/src/testing_utils.rs b/p2p/src/testing_utils.rs
index 03c8987ff..deec1b049 100644
--- a/p2p/src/testing_utils.rs
+++ b/p2p/src/testing_utils.rs
@@ -18,9 +18,14 @@
 use std::{
     fmt::Debug,
     net::{IpAddr, Ipv6Addr, SocketAddr},
+    sync::{
+        atomic::{AtomicU64, Ordering},
+        Arc,
+    },
     time::Duration,
 };
 
+use common::time_getter::TimeGetter;
 use crypto::random::{make_pseudo_rng, Rng};
 use tokio::time::timeout;
 
@@ -224,3 +229,36 @@ pub fn peerdb_inmemory_store() -> PeerDbStorageImpl<storage::inmemory::InMemory>
     let storage = storage::inmemory::InMemory::new();
     PeerDbStorageImpl::new(storage).unwrap()
 }
+
+pub struct P2pTestTimeGetter {
+    current_time_millis: Arc<AtomicU64>,
+}
+
+impl P2pTestTimeGetter {
+    pub fn new() -> Self {
+        let current_time = std::time::SystemTime::now()
+            .duration_since(std::time::SystemTime::UNIX_EPOCH)
+            .unwrap();
+        let current_time_millis = Arc::new(AtomicU64::new(current_time.as_millis() as u64));
+        Self {
+            current_time_millis,
+        }
+    }
+
+    pub fn get_time_getter(&self) -> TimeGetter {
+        let current_time = Arc::clone(&self.current_time_millis);
+        TimeGetter::new(Arc::new(move || {
+            Duration::from_millis(current_time.load(Ordering::SeqCst))
+        }))
+    }
+
+    pub async fn advance_time(&self, duration: Duration) {
+        tokio::time::pause();
+        self.current_time_millis.store(
+            self.current_time_millis.load(Ordering::SeqCst) + duration.as_millis() as u64,
+            Ordering::SeqCst,
+        );
+        tokio::time::advance(duration).await;
+        tokio::time::resume();
+    }
+}