use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::Notify;
use crate::bus::Subscription;
use crate::event::Event;
use crate::EventBus;
pub struct EventCollector<E: Event> {
events: Arc<Mutex<Vec<E>>>,
notify: Arc<Notify>,
_subscription: Subscription,
}
impl<E: Event> EventCollector<E> {
pub async fn new(bus: &EventBus) -> Self {
let events: Arc<Mutex<Vec<E>>> = Arc::new(Mutex::new(Vec::new()));
let notify = Arc::new(Notify::new());
let events_clone = events.clone();
let notify_clone = notify.clone();
let subscription = bus
.subscribe(move |event: E| {
let events = events_clone.clone();
let notify = notify_clone.clone();
async move {
events.lock().unwrap().push(event);
notify.notify_one();
Ok(())
}
})
.await
.expect("EventCollector::new() subscribe failed");
Self {
events,
notify,
_subscription: subscription,
}
}
pub fn collect_now(&self) -> Vec<E> {
self.events.lock().unwrap().clone()
}
pub async fn wait_for(&self, count: usize, timeout: Duration) -> Vec<E> {
let deadline = tokio::time::Instant::now() + timeout;
loop {
{
let events = self.events.lock().unwrap();
if events.len() >= count {
return events.clone();
}
}
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
if remaining.is_zero() {
return self.collect_now();
}
tokio::select! {
_ = self.notify.notified() => {
}
_ = tokio::time::sleep(remaining) => {
return self.collect_now();
}
}
}
}
pub fn assert_count(&self, expected: usize) {
let events = self.collect_now();
assert_eq!(
events.len(),
expected,
"EventCollector<{}>: 期望 {} 个事件,实际收集到 {} 个",
std::any::type_name::<E>(),
expected,
events.len()
);
}
pub fn assert_contains(&self, predicate: impl Fn(&E) -> bool) {
let events = self.collect_now();
let found = events.iter().any(|e| predicate(e));
assert!(
found,
"EventCollector<{}>: 未找到满足条件的事件(共 {} 个事件)",
std::any::type_name::<E>(),
events.len()
);
}
pub fn assert_not_contains(&self, predicate: impl Fn(&E) -> bool) {
let events = self.collect_now();
let found = events.iter().any(|e| predicate(e));
assert!(
!found,
"EventCollector<{}>: 意外找到了满足条件的事件(共 {} 个事件)",
std::any::type_name::<E>(),
events.len()
);
}
}
#[cfg(test)]
mod tests {
use serde::{Deserialize, Serialize};
use crate::event::Event;
use crate::testing::EventCollector;
use crate::EventBus;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct TestEvent {
id: u64,
name: String,
}
impl Event for TestEvent {
fn event_name() -> &'static str {
"test.event"
}
}
#[tokio::test]
async fn test_collect_now() {
let bus = EventBus::new();
let collector = EventCollector::<TestEvent>::new(&bus).await;
bus.publish(TestEvent {
id: 1,
name: "first".into(),
})
.await
.unwrap();
bus.publish(TestEvent {
id: 2,
name: "second".into(),
})
.await
.unwrap();
let events = collector.wait_for(2, std::time::Duration::from_secs(2)).await;
assert_eq!(events.len(), 2);
assert_eq!(events[0].id, 1);
assert_eq!(events[0].name, "first");
assert_eq!(events[1].id, 2);
assert_eq!(events[1].name, "second");
}
#[tokio::test]
async fn test_wait_for_success() {
let bus = EventBus::new();
let collector = EventCollector::<TestEvent>::new(&bus).await;
let bus_clone = bus.clone();
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
bus_clone
.publish(TestEvent {
id: 1,
name: "delayed".into(),
})
.await
.unwrap();
});
let events = collector
.wait_for(1, std::time::Duration::from_secs(2))
.await;
assert_eq!(events.len(), 1);
assert_eq!(events[0].name, "delayed");
}
#[tokio::test]
async fn test_wait_for_timeout() {
let bus = EventBus::new();
let collector = EventCollector::<TestEvent>::new(&bus).await;
let events = collector
.wait_for(5, std::time::Duration::from_millis(100))
.await;
assert_eq!(events.len(), 0);
}
#[tokio::test]
async fn test_assert_count() {
let bus = EventBus::new();
let collector = EventCollector::<TestEvent>::new(&bus).await;
bus.publish(TestEvent {
id: 1,
name: "a".into(),
})
.await
.unwrap();
bus.publish(TestEvent {
id: 2,
name: "b".into(),
})
.await
.unwrap();
let _ = collector
.wait_for(2, std::time::Duration::from_secs(2))
.await;
collector.assert_count(2);
}
#[tokio::test]
async fn test_assert_contains() {
let bus = EventBus::new();
let collector = EventCollector::<TestEvent>::new(&bus).await;
bus.publish(TestEvent {
id: 1,
name: "alice".into(),
})
.await
.unwrap();
let _ = collector
.wait_for(1, std::time::Duration::from_secs(2))
.await;
collector.assert_contains(|e| e.name == "alice");
collector.assert_not_contains(|e| e.name == "bob");
}
#[tokio::test]
async fn test_multiple_collectors() {
let bus = EventBus::new();
let collector1 = EventCollector::<TestEvent>::new(&bus).await;
let collector2 = EventCollector::<TestEvent>::new(&bus).await;
bus.publish(TestEvent {
id: 1,
name: "shared".into(),
})
.await
.unwrap();
let events1 = collector1
.wait_for(1, std::time::Duration::from_secs(2))
.await;
let events2 = collector2
.wait_for(1, std::time::Duration::from_secs(2))
.await;
assert_eq!(events1.len(), 1);
assert_eq!(events2.len(), 1);
assert_eq!(events1[0].name, "shared");
assert_eq!(events2[0].name, "shared");
}
}