Skip to main content

af_notify/
testing.rs

1//! In-memory notification store and senders for consumer tests.
2
3use std::collections::BTreeMap;
4use std::sync::Mutex;
5use std::time::Duration;
6
7use af_context::{NotificationId, RequestContext};
8use async_trait::async_trait;
9use chrono::Utc;
10
11use crate::{
12    AttemptEvent, DeliveryClass, DeliveryContext, DurableOutbox, NewOutboxItem, Notification,
13    NotificationStatus, NotifyError, NotifyStoreError, OutboxItem, Sender,
14};
15
16/// In-memory durable queue with lease fencing.
17pub struct MemoryOutbox {
18    items: Mutex<BTreeMap<NotificationId, OutboxItem>>,
19    idempotency: Mutex<BTreeMap<(String, String, String), NotificationId>>,
20    changes: tokio::sync::broadcast::Sender<NotificationId>,
21}
22
23impl Default for MemoryOutbox {
24    fn default() -> Self {
25        let (changes, _) = tokio::sync::broadcast::channel(32);
26        Self {
27            items: Mutex::new(BTreeMap::new()),
28            idempotency: Mutex::new(BTreeMap::new()),
29            changes,
30        }
31    }
32}
33
34#[async_trait]
35impl DurableOutbox for MemoryOutbox {
36    fn subscribe(&self) -> tokio::sync::broadcast::Receiver<NotificationId> {
37        self.changes.subscribe()
38    }
39    async fn enqueue(&self, item: NewOutboxItem) -> Result<OutboxItem, NotifyStoreError> {
40        let key = (
41            item.tenant_id.to_string(),
42            item.subject_id.to_string(),
43            item.idempotency_key.clone(),
44        );
45        let mut idempotency = self
46            .idempotency
47            .lock()
48            .unwrap_or_else(std::sync::PoisonError::into_inner);
49        let mut items = self
50            .items
51            .lock()
52            .unwrap_or_else(std::sync::PoisonError::into_inner);
53        if let Some(existing) = idempotency.get(&key).and_then(|id| items.get(id)) {
54            return if existing.channel == item.channel
55                && existing.recipient == item.recipient
56                && existing.notification == item.notification
57                && existing.max_attempts == item.max_attempts.max(1)
58            {
59                Ok(existing.clone())
60            } else {
61                Err(NotifyStoreError::IdempotencyConflict)
62            };
63        }
64        let id = NotificationId::parse(format!("notification-{}", items.len() + 1))
65            .map_err(|error| NotifyStoreError::Invalid(error.to_string()))?;
66        let now = Utc::now();
67        let record = OutboxItem {
68            id: id.clone(),
69            tenant_id: item.tenant_id,
70            subject_id: item.subject_id,
71            channel: item.channel,
72            recipient: item.recipient,
73            notification: item.notification,
74            status: NotificationStatus::Pending,
75            attempts: 0,
76            max_attempts: item.max_attempts.max(1),
77            lease_version: 0,
78            last_error: None,
79            created_at: now,
80            updated_at: now,
81        };
82        items.insert(id.clone(), record.clone());
83        idempotency.insert(key, id);
84        let _ = self.changes.send(record.id.clone());
85        Ok(record)
86    }
87    async fn get(
88        &self,
89        context: &RequestContext,
90        id: &NotificationId,
91    ) -> Result<OutboxItem, NotifyStoreError> {
92        self.items
93            .lock()
94            .unwrap_or_else(std::sync::PoisonError::into_inner)
95            .get(id)
96            .filter(|item| {
97                item.tenant_id == context.tenant_id && item.subject_id == context.subject_id
98            })
99            .cloned()
100            .ok_or(NotifyStoreError::NotFound)
101    }
102    async fn list(
103        &self,
104        context: &RequestContext,
105        status: Option<NotificationStatus>,
106        limit: usize,
107        after: Option<&NotificationId>,
108    ) -> Result<Vec<OutboxItem>, NotifyStoreError> {
109        Ok(self
110            .items
111            .lock()
112            .unwrap_or_else(std::sync::PoisonError::into_inner)
113            .values()
114            .filter(|item| {
115                item.tenant_id == context.tenant_id
116                    && item.subject_id == context.subject_id
117                    && status.is_none_or(|status| item.status == status)
118                    && after.is_none_or(|after| item.id > *after)
119            })
120            .take(limit.clamp(1, 101))
121            .cloned()
122            .collect())
123    }
124    async fn attempts(
125        &self,
126        context: &RequestContext,
127        id: &NotificationId,
128    ) -> Result<Vec<AttemptEvent>, NotifyStoreError> {
129        self.get(context, id).await.map(|_| Vec::new())
130    }
131    async fn claim(
132        &self,
133        _worker_id: &str,
134        _lease_secs: i64,
135        batch: usize,
136    ) -> Result<Vec<OutboxItem>, NotifyStoreError> {
137        let mut items = self
138            .items
139            .lock()
140            .unwrap_or_else(std::sync::PoisonError::into_inner);
141        let ids: Vec<_> = items
142            .values()
143            .filter(|item| {
144                matches!(
145                    item.status,
146                    NotificationStatus::Pending | NotificationStatus::RetryScheduled
147                )
148            })
149            .take(batch.clamp(1, 100))
150            .map(|item| item.id.clone())
151            .collect();
152        let claimed: Vec<_> = ids
153            .into_iter()
154            .filter_map(|id| {
155                let item = items.get_mut(&id)?;
156                item.status = NotificationStatus::Sending;
157                item.attempts += 1;
158                item.lease_version += 1;
159                Some(item.clone())
160            })
161            .collect();
162        for item in &claimed {
163            let _ = self.changes.send(item.id.clone());
164        }
165        Ok(claimed)
166    }
167    async fn mark_sent(
168        &self,
169        id: &NotificationId,
170        lease_version: i64,
171    ) -> Result<(), NotifyStoreError> {
172        let mut items = self
173            .items
174            .lock()
175            .unwrap_or_else(std::sync::PoisonError::into_inner);
176        let item = items.get_mut(id).ok_or(NotifyStoreError::NotFound)?;
177        if item.lease_version != lease_version {
178            return Err(NotifyStoreError::LeaseLost);
179        }
180        item.status = NotificationStatus::Sent;
181        let _ = self.changes.send(id.clone());
182        Ok(())
183    }
184    async fn record_failure(
185        &self,
186        id: &NotificationId,
187        lease_version: i64,
188        class: DeliveryClass,
189        _delay: Duration,
190    ) -> Result<NotificationStatus, NotifyStoreError> {
191        let mut items = self
192            .items
193            .lock()
194            .unwrap_or_else(std::sync::PoisonError::into_inner);
195        let item = items.get_mut(id).ok_or(NotifyStoreError::NotFound)?;
196        if item.lease_version != lease_version {
197            return Err(NotifyStoreError::LeaseLost);
198        }
199        item.status = if class == DeliveryClass::Permanent || item.attempts >= item.max_attempts {
200            NotificationStatus::DeadLetter
201        } else {
202            NotificationStatus::RetryScheduled
203        };
204        let _ = self.changes.send(id.clone());
205        Ok(item.status)
206    }
207    async fn retry_dead_letter(
208        &self,
209        context: &RequestContext,
210        id: &NotificationId,
211    ) -> Result<OutboxItem, NotifyStoreError> {
212        let mut items = self
213            .items
214            .lock()
215            .unwrap_or_else(std::sync::PoisonError::into_inner);
216        let item = items
217            .get_mut(id)
218            .filter(|item| item.tenant_id == context.tenant_id)
219            .ok_or(NotifyStoreError::NotFound)?;
220        if item.status != NotificationStatus::DeadLetter {
221            return Err(NotifyStoreError::Invalid(
222                "notification is not dead-lettered".into(),
223            ));
224        }
225        item.status = NotificationStatus::Pending;
226        item.attempts = 0;
227        let _ = self.changes.send(id.clone());
228        Ok(item.clone())
229    }
230}
231
232/// Sender that captures rendered deliveries.
233#[derive(Default)]
234pub struct CapturingSender {
235    /// Captured recipient and body pairs.
236    pub deliveries: Mutex<Vec<(String, String)>>,
237}
238
239#[async_trait]
240impl Sender for CapturingSender {
241    fn name(&self) -> &str {
242        "capture"
243    }
244    async fn send(
245        &self,
246        context: &DeliveryContext,
247        recipient: &str,
248        notification: &Notification,
249    ) -> Result<(), NotifyError> {
250        if context.cancellation.is_cancelled() {
251            return Err(NotifyError::Cancelled);
252        }
253        self.deliveries
254            .lock()
255            .unwrap_or_else(std::sync::PoisonError::into_inner)
256            .push((recipient.to_owned(), notification.to_plain_text()));
257        Ok(())
258    }
259}
260
261/// Sender that always returns one configured failure.
262pub struct FailingSender(pub NotifyError);
263
264#[async_trait]
265impl Sender for FailingSender {
266    fn name(&self) -> &str {
267        "fail"
268    }
269    async fn send(
270        &self,
271        _context: &DeliveryContext,
272        _recipient: &str,
273        _notification: &Notification,
274    ) -> Result<(), NotifyError> {
275        Err(self.0.clone())
276    }
277}