zebra_network/address_book_peers/
mock.rs

1//! Mock [`AddressBookPeers`] for use in tests.
2
3use std::{
4    sync::{Arc, Mutex},
5    time::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        self.recently_live_peers.push(
29            MetaAddr::new_responded(peer).into_new_meta_addr(
30                Instant::now(),
31                chrono::Utc::now()
32                    .try_into()
33                    .expect("will succeed until 2038"),
34            ),
35        );
36        true
37    }
38}
39
40impl AddressBookPeers for MockAddressBookPeers {
41    fn recently_live_peers(&self, now: chrono::DateTime<chrono::Utc>) -> Vec<MetaAddr> {
42        self.recently_live_peers
43            .iter()
44            .filter(|peer| peer.was_recently_live(now))
45            .cloned()
46            .collect()
47    }
48
49    fn add_peer(&mut self, peer: PeerSocketAddr) -> bool {
50        self.add_peer(peer)
51    }
52}
53
54/// A shared version of [`MockAddressBookPeers`] that can be used in shared state tests.
55pub struct SharedMockAddressBookPeers {
56    inner: Arc<Mutex<MockAddressBookPeers>>,
57}
58
59impl SharedMockAddressBookPeers {
60    /// Creates a new instance of [`SharedMockAddressBookPeers`].
61    pub fn new() -> Self {
62        Self {
63            inner: Arc::new(Mutex::new(MockAddressBookPeers::default())),
64        }
65    }
66}
67
68impl Default for SharedMockAddressBookPeers {
69    fn default() -> Self {
70        Self::new()
71    }
72}
73
74impl AddressBookPeers for SharedMockAddressBookPeers {
75    fn recently_live_peers(&self, now: chrono::DateTime<chrono::Utc>) -> Vec<MetaAddr> {
76        self.inner.lock().unwrap().recently_live_peers(now)
77    }
78
79    fn add_peer(&mut self, peer: PeerSocketAddr) -> bool {
80        self.inner.lock().unwrap().add_peer(peer)
81    }
82}
83
84impl Clone for SharedMockAddressBookPeers {
85    fn clone(&self) -> Self {
86        Self {
87            inner: self.inner.clone(),
88        }
89    }
90}