Skip to main content

zebra_network/address_book_peers/
mock.rs

1//! Mock [`AddressBookPeers`] for use in tests.
2
3use std::{
4    sync::{Arc, Mutex},
5    time::{Duration, Instant},
6};
7
8use crate::{meta_addr::MetaAddr, AddressBookPeers, PeerSocketAddr};
9
10/// A mock [`AddressBookPeers`] implementation that's always empty.
11#[derive(Debug, Default, Clone)]
12pub struct MockAddressBookPeers {
13    /// Return value for mock `recently_live_peers` method.
14    recently_live_peers: Vec<MetaAddr>,
15}
16
17impl MockAddressBookPeers {
18    /// Creates a new [`MockAddressBookPeers`]
19    pub fn new(recently_live_peers: Vec<MetaAddr>) -> Self {
20        Self {
21            recently_live_peers,
22        }
23    }
24
25    /// Adds a peer to the mock address book.
26    pub fn add_peer(&mut self, peer: PeerSocketAddr) -> bool {
27        // The real add peer will use `MetaAddr::new_initial_peer` but we just want to get a `MetaAddr` for the mock.
28        let rtt = Duration::from_millis(100);
29        self.recently_live_peers.push(
30            MetaAddr::new_responded(peer, Some(rtt)).into_new_meta_addr(
31                Instant::now(),
32                chrono::Utc::now()
33                    .try_into()
34                    .expect("will succeed until 2038"),
35            ),
36        );
37        true
38    }
39}
40
41impl AddressBookPeers for MockAddressBookPeers {
42    fn recently_live_peers(&self, now: chrono::DateTime<chrono::Utc>) -> Vec<MetaAddr> {
43        self.recently_live_peers
44            .iter()
45            .filter(|peer| peer.was_recently_live(now))
46            .cloned()
47            .collect()
48    }
49
50    fn add_peer(&mut self, peer: PeerSocketAddr) -> bool {
51        self.add_peer(peer)
52    }
53}
54
55/// A shared version of [`MockAddressBookPeers`] that can be used in shared state tests.
56pub struct SharedMockAddressBookPeers {
57    inner: Arc<Mutex<MockAddressBookPeers>>,
58}
59
60impl SharedMockAddressBookPeers {
61    /// Creates a new instance of [`SharedMockAddressBookPeers`].
62    pub fn new() -> Self {
63        Self {
64            inner: Arc::new(Mutex::new(MockAddressBookPeers::default())),
65        }
66    }
67}
68
69impl Default for SharedMockAddressBookPeers {
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75impl AddressBookPeers for SharedMockAddressBookPeers {
76    fn recently_live_peers(&self, now: chrono::DateTime<chrono::Utc>) -> Vec<MetaAddr> {
77        self.inner.lock().unwrap().recently_live_peers(now)
78    }
79
80    fn add_peer(&mut self, peer: PeerSocketAddr) -> bool {
81        self.inner.lock().unwrap().add_peer(peer)
82    }
83}
84
85impl Clone for SharedMockAddressBookPeers {
86    fn clone(&self) -> Self {
87        Self {
88            inner: self.inner.clone(),
89        }
90    }
91}