af_notify/

testing.rs

//! In-memory notification store and senders for consumer tests.

use std::collections::BTreeMap;
use std::sync::Mutex;
use std::time::Duration;

use af_context::{NotificationId, RequestContext};
use async_trait::async_trait;
use chrono::Utc;

use crate::{
    AttemptEvent, DeliveryClass, DeliveryContext, DurableOutbox, NewOutboxItem, Notification,
    NotificationStatus, NotifyError, NotifyStoreError, OutboxItem, Sender,
};

/// In-memory durable queue with lease fencing.
pub struct MemoryOutbox {
    notifications: bool,
    items: Mutex<BTreeMap<NotificationId, OutboxItem>>,
    idempotency: Mutex<BTreeMap<(String, String, String), NotificationId>>,
    changes: tokio::sync::broadcast::Sender<NotificationId>,
}

impl Default for MemoryOutbox {
    fn default() -> Self {
        let (changes, _) = tokio::sync::broadcast::channel(32);
        Self {
            notifications: true,
            items: Mutex::new(BTreeMap::new()),
            idempotency: Mutex::new(BTreeMap::new()),
            changes,
        }
    }
}

impl MemoryOutbox {
    /// Model a worker in another process: durable reads work without local broadcasts.
    pub fn without_notifications() -> Self {
        Self {
            notifications: false,
            ..Self::default()
        }
    }
}

#[async_trait]
impl DurableOutbox for MemoryOutbox {
    fn subscribe(&self) -> tokio::sync::broadcast::Receiver<NotificationId> {
        if self.notifications {
            self.changes.subscribe()
        } else {
            let (_, receiver) = tokio::sync::broadcast::channel(1);
            receiver
        }
    }
    async fn enqueue(&self, item: NewOutboxItem) -> Result<OutboxItem, NotifyStoreError> {
        let key = (
            item.tenant_id.to_string(),
            item.subject_id.to_string(),
            item.idempotency_key.clone(),
        );
        let mut idempotency = self
            .idempotency
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut items = self
            .items
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(existing) = idempotency.get(&key).and_then(|id| items.get(id)) {
            return if existing.channel == item.channel
                && existing.recipient == item.recipient
                && existing.notification == item.notification
                && existing.max_attempts == item.max_attempts.max(1)
            {
                Ok(existing.clone())
            } else {
                Err(NotifyStoreError::IdempotencyConflict)
            };
        }
        let id = NotificationId::parse(format!("notification-{}", items.len() + 1))
            .map_err(|error| NotifyStoreError::Invalid(error.to_string()))?;
        let now = Utc::now();
        let record = OutboxItem {
            id: id.clone(),
            tenant_id: item.tenant_id,
            subject_id: item.subject_id,
            channel: item.channel,
            recipient: item.recipient,
            notification: item.notification,
            status: NotificationStatus::Pending,
            attempts: 0,
            max_attempts: item.max_attempts.max(1),
            lease_version: 0,
            last_error: None,
            created_at: now,
            updated_at: now,
        };
        items.insert(id.clone(), record.clone());
        idempotency.insert(key, id);
        let _ = self.changes.send(record.id.clone());
        Ok(record)
    }
    async fn get(
        &self,
        context: &RequestContext,
        id: &NotificationId,
    ) -> Result<OutboxItem, NotifyStoreError> {
        self.items
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .get(id)
            .filter(|item| {
                item.tenant_id == context.tenant_id && item.subject_id == context.subject_id
            })
            .cloned()
            .ok_or(NotifyStoreError::NotFound)
    }
    async fn list(
        &self,
        context: &RequestContext,
        status: Option<NotificationStatus>,
        limit: usize,
        after: Option<&NotificationId>,
    ) -> Result<Vec<OutboxItem>, NotifyStoreError> {
        Ok(self
            .items
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .values()
            .filter(|item| {
                item.tenant_id == context.tenant_id
                    && item.subject_id == context.subject_id
                    && status.is_none_or(|status| item.status == status)
                    && after.is_none_or(|after| item.id > *after)
            })
            .take(limit.clamp(1, 101))
            .cloned()
            .collect())
    }
    async fn attempts(
        &self,
        context: &RequestContext,
        id: &NotificationId,
    ) -> Result<Vec<AttemptEvent>, NotifyStoreError> {
        self.get(context, id).await.map(|_| Vec::new())
    }
    async fn claim(
        &self,
        _worker_id: &str,
        _lease_secs: i64,
        batch: usize,
    ) -> Result<Vec<OutboxItem>, NotifyStoreError> {
        let mut items = self
            .items
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let ids: Vec<_> = items
            .values()
            .filter(|item| {
                matches!(
                    item.status,
                    NotificationStatus::Pending | NotificationStatus::RetryScheduled
                )
            })
            .take(batch.clamp(1, 100))
            .map(|item| item.id.clone())
            .collect();
        let claimed: Vec<_> = ids
            .into_iter()
            .filter_map(|id| {
                let item = items.get_mut(&id)?;
                item.status = NotificationStatus::Sending;
                item.attempts += 1;
                item.lease_version += 1;
                Some(item.clone())
            })
            .collect();
        for item in &claimed {
            let _ = self.changes.send(item.id.clone());
        }
        Ok(claimed)
    }
    async fn mark_sent(
        &self,
        id: &NotificationId,
        lease_version: i64,
    ) -> Result<(), NotifyStoreError> {
        let mut items = self
            .items
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let item = items.get_mut(id).ok_or(NotifyStoreError::NotFound)?;
        if item.lease_version != lease_version {
            return Err(NotifyStoreError::LeaseLost);
        }
        item.status = NotificationStatus::Sent;
        let _ = self.changes.send(id.clone());
        Ok(())
    }
    async fn record_failure(
        &self,
        id: &NotificationId,
        lease_version: i64,
        class: DeliveryClass,
        _delay: Duration,
    ) -> Result<NotificationStatus, NotifyStoreError> {
        let mut items = self
            .items
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let item = items.get_mut(id).ok_or(NotifyStoreError::NotFound)?;
        if item.lease_version != lease_version {
            return Err(NotifyStoreError::LeaseLost);
        }
        item.status = if class == DeliveryClass::Permanent || item.attempts >= item.max_attempts {
            NotificationStatus::DeadLetter
        } else {
            NotificationStatus::RetryScheduled
        };
        let _ = self.changes.send(id.clone());
        Ok(item.status)
    }
    async fn retry_dead_letter(
        &self,
        context: &RequestContext,
        id: &NotificationId,
    ) -> Result<OutboxItem, NotifyStoreError> {
        let mut items = self
            .items
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let item = items
            .get_mut(id)
            .filter(|item| item.tenant_id == context.tenant_id)
            .ok_or(NotifyStoreError::NotFound)?;
        if item.status != NotificationStatus::DeadLetter {
            return Err(NotifyStoreError::Invalid(
                "notification is not dead-lettered".into(),
            ));
        }
        item.status = NotificationStatus::Pending;
        item.attempts = 0;
        let _ = self.changes.send(id.clone());
        Ok(item.clone())
    }
}

/// Sender that captures rendered deliveries.
#[derive(Default)]
pub struct CapturingSender {
    /// Captured recipient and body pairs.
    pub deliveries: Mutex<Vec<(String, String)>>,
}

#[async_trait]
impl Sender for CapturingSender {
    fn name(&self) -> &str {
        "capture"
    }
    async fn send(
        &self,
        context: &DeliveryContext,
        recipient: &str,
        notification: &Notification,
    ) -> Result<(), NotifyError> {
        if context.cancellation.is_cancelled() {
            return Err(NotifyError::Cancelled);
        }
        self.deliveries
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .push((recipient.to_owned(), notification.to_plain_text()));
        Ok(())
    }
}

/// Sender that always returns one configured failure.
pub struct FailingSender(pub NotifyError);

#[async_trait]
impl Sender for FailingSender {
    fn name(&self) -> &str {
        "fail"
    }
    async fn send(
        &self,
        _context: &DeliveryContext,
        _recipient: &str,
        _notification: &Notification,
    ) -> Result<(), NotifyError> {
        Err(self.0.clone())
    }
}