test(apollo_network): unit tests for the EventWakerManager struct

Author: guyf-starkware

Cargo.lock

@@ -310,6 +310,7 @@ unsigned-varint = "0.8.0"
 url = "2.5.0"
 validator = "0.12"
 void = "1.0.2"
+waker-fn = "1"
 workspace_tests.path = "workspace_tests"
 zstd = "0.13.1"
 

Cargo.toml

@@ -54,6 +54,7 @@ rstest.workspace = true
 tokio = { workspace = true, features = ["full", "sync", "test-util"] }
 tokio-stream.workspace = true
 void.workspace = true
+waker-fn.workspace = true
 
 [lints]
 workspace = true

crates/apollo_network/Cargo.toml

@@ -0,0 +1,40 @@
+use super::EventWakerManager;
+use crate::discovery::behaviours::test_util::MockWakerWrapper;
+
+#[test]
+fn no_wakers_added() {
+    let mut waker_manager = EventWakerManager::default();
+    waker_manager.wake();
+}
+
+#[test]
+fn wakes_all_wakers() {
+    let waker_wrapper1 = MockWakerWrapper::new();
+    let waker_wrapper2 = MockWakerWrapper::new();
+
+    let mut waker_manager = EventWakerManager::default();
+    waker_manager.add_waker(waker_wrapper1.get_waker());
+    waker_manager.add_waker(waker_wrapper2.get_waker());
+
+    assert_eq!(waker_wrapper1.times_woken(), 0);
+    assert_eq!(waker_wrapper2.times_woken(), 0);
+
+    waker_manager.wake();
+
+    assert_eq!(waker_wrapper1.times_woken(), 1);
+    assert_eq!(waker_wrapper2.times_woken(), 1);
+}
+
+#[test]
+fn does_not_wake_waker_twice() {
+    let waker_wrapper = MockWakerWrapper::new();
+
+    let mut waker_manager = EventWakerManager::default();
+    waker_manager.add_waker(waker_wrapper.get_waker());
+
+    waker_manager.wake();
+    assert_eq!(waker_wrapper.times_woken(), 1);
+
+    waker_manager.wake();
+    assert_eq!(waker_wrapper.times_woken(), 1);
+}

crates/apollo_network/src/discovery/behaviours/event_waker_manager_test.rs

@@ -1,7 +1,11 @@
 pub mod bootstrapping;
 mod event_waker_manager;
 pub mod kad_requesting;
+mod test_util;
 mod time_waker_manager;
 
+#[cfg(test)]
+mod event_waker_manager_test;
+
 use event_waker_manager::EventWakerManager;
 use time_waker_manager::TimeWakerManager;

crates/apollo_network/src/discovery/behaviours/mod.rs

@@ -0,0 +1,31 @@
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::task::Waker;
+
+use waker_fn::waker_fn;
+
+#[cfg(test)]
+pub(super) struct MockWakerWrapper {
+    times_woken: Arc<AtomicUsize>,
+    waker: Waker,
+}
+
+#[cfg(test)]
+impl MockWakerWrapper {
+    pub fn new() -> Self {
+        let times_woken = Arc::new(AtomicUsize::new(0));
+        let times_woken_clone = times_woken.clone();
+        let waker = waker_fn(move || {
+            times_woken_clone.fetch_add(1, Ordering::SeqCst);
+        });
+        Self { times_woken, waker }
+    }
+
+    pub fn get_waker(&self) -> &Waker {
+        &self.waker
+    }
+
+    pub fn times_woken(&self) -> usize {
+        self.times_woken.load(Ordering::SeqCst)
+    }
+}