1 files changed, 72 insertions, 17 deletions
diff --git a/src/test/webpush.rs b/src/test/webpush.rs
index c86d03f..55caf19 100644
--- a/src/test/webpush.rs
+++ b/src/test/webpush.rs
@@ -1,37 +1,92 @@
 use std::{
+    any::Any,
+    collections::{HashMap, HashSet},
     mem,
-    sync::{Arc, Mutex},
+    sync::{Arc, Mutex, MutexGuard},
 };
 
-use web_push::{WebPushClient, WebPushError, WebPushMessage};
+use web_push::{PartialVapidSignatureBuilder, SubscriptionInfo, WebPushError};
+
+use crate::{error::failed::Failed, push::Publish};
 
 #[derive(Clone)]
-pub struct Client {
-    sent: Arc<Mutex<Vec<WebPushMessage>>>,
+pub struct Client(Arc<Mutex<ClientInner>>);
+
+#[derive(Default)]
+struct ClientInner {
+    sent: Vec<Publication>,
+    planned_failures: HashMap<SubscriptionInfo, WebPushError>,
 }
 
 impl Client {
     pub fn new() -> Self {
-        Self {
-            sent: Arc::default(),
-        }
+        Self(Arc::default())
+    }
+
+    fn inner(&self) -> MutexGuard<'_, ClientInner> {
+        self.0.lock().unwrap()
     }
 
     // Clears the list of sent messages (for all clones of this Client) when called, because we
-    // can't clone `WebPushMessage`s so we either need to move them or try to reconstruct them,
-    // either of which sucks but moving them sucks less.
-    pub fn sent(&self) -> Vec<WebPushMessage> {
-        let mut sent = self.sent.lock().unwrap();
-        mem::replace(&mut *sent, Vec::new())
+    // can't clone `Publications`s, so we either need to move them or try to reconstruct them.
+    pub fn sent(&self) -> Vec<Publication> {
+        let sent = &mut self.inner().sent;
+        mem::take(sent)
+    }
+
+    pub fn fail_next(&self, subscription_info: &SubscriptionInfo, err: WebPushError) {
+        let planned_failures = &mut self.inner().planned_failures;
+        planned_failures.insert(subscription_info.clone(), err);
     }
 }
 
 #[async_trait::async_trait]
-impl WebPushClient for Client {
-    async fn send(&self, message: WebPushMessage) -> Result<(), WebPushError> {
-        let mut sent = self.sent.lock().unwrap();
-        sent.push(message);
+impl Publish for Client {
+    async fn publish<'s, M>(
+        &self,
+        message: M,
+        _: &PartialVapidSignatureBuilder,
+        subscriptions: impl IntoIterator<Item = &'s SubscriptionInfo> + Send,
+    ) -> Result<Vec<(&'s SubscriptionInfo, WebPushError)>, Failed>
+    where
+        M: Send + 'static,
+    {
+        let mut inner = self.inner();
+        let message: Box<dyn Any + Send> = Box::new(message);
+
+        let mut recipients = HashSet::new();
+        let mut failures = Vec::new();
+        for subscription in subscriptions {
+            recipients.insert(subscription.clone());
+            if let Some(err) = inner.planned_failures.remove(subscription) {
+                failures.push((subscription, err));
+            }
+        }
+
+        let publication = Publication {
+            message,
+            recipients,
+        };
+        inner.sent.push(publication);
+
+        Ok(failures)
+    }
+}
+
+#[derive(Debug)]
+pub struct Publication {
+    pub message: Box<dyn Any + Send>,
+    pub recipients: HashSet<SubscriptionInfo>,
+}
 
-        Ok(())
+impl Publication {
+    pub fn message_eq<M>(&self, candidate: &M) -> bool
+    where
+        M: PartialEq + 'static,
+    {
+        match self.message.downcast_ref::<M>() {
+            None => false,
+            Some(message) => message == candidate,
+        }
     }
 }