use crate::Span;
use commonware_utils::channel::{fallible::FallibleExt, mpsc};
use std::collections::HashMap;
#[derive(Debug)]
pub enum Event<K, V> {
Success(K, V),
Failed(K),
}
#[derive(Clone)]
pub struct Consumer<K: Span, V> {
sender: mpsc::UnboundedSender<Event<K, V>>,
expected: HashMap<K, V>,
}
impl<K: Span, V: Clone + PartialEq> Consumer<K, V> {
pub fn new() -> (Self, mpsc::UnboundedReceiver<Event<K, V>>) {
let (sender, receiver) = mpsc::unbounded_channel();
(
Self {
sender,
expected: HashMap::new(),
},
receiver,
)
}
pub fn dummy() -> Self {
let (sender, _) = mpsc::unbounded_channel();
Self {
sender,
expected: HashMap::new(),
}
}
pub fn add_expected(&mut self, k: K, v: V) {
self.expected.insert(k, v);
}
pub fn pop_expected(&mut self, k: &K) -> Option<V> {
self.expected.remove(k)
}
}
impl<K: Span, V: Clone + PartialEq + Send + 'static> crate::Consumer for Consumer<K, V> {
type Key = K;
type Value = V;
type Failure = ();
async fn deliver(&mut self, key: Self::Key, value: Self::Value) -> bool {
let valid = self.expected.get(&key).is_none_or(|v| v == &value);
if valid {
self.sender.send_lossy(Event::Success(key, value));
}
valid
}
async fn failed(&mut self, key: Self::Key, _failure: ()) {
self.sender.send_lossy(Event::Failed(key));
}
}