azums_core/backend/
mock.rs

1use crate::{
2    backend::{
3        memory::MemoryBackend, JobExplanation, NotificationStream, ObservabilityBackend,
4        QueueMetrics, StorageBackend, StreamBackend,
5    },
6    model::{ConsumerGroupStatus, Event, Job, JobListItem, NewEvent, NewJob},
7};
8use async_trait::async_trait;
9use chrono::{DateTime, Utc};
10use std::sync::{Arc, Mutex};
11use uuid::Uuid;
12
13/// Log record representing a single call executed against a [`MockBackend`].
14#[derive(Debug, Clone)]
15pub enum CallRecord {
16    RunMigrations,
17    HealthCheck,
18    Enqueue(NewJob),
19    Subscribe(String),
20    PublishStream {
21        stream: String,
22        event: NewEvent,
23    },
24    SubscribeStream {
25        stream: String,
26        consumer_group: String,
27        last_seq: Option<i64>,
28    },
29    AckStream {
30        stream: String,
31        consumer_group: String,
32        seq: i64,
33    },
34    ReadEventsStream {
35        stream: String,
36        after_seq: i64,
37        limit: i64,
38    },
39    PruneEventsStream {
40        stream: String,
41        through_seq: i64,
42    },
43    ConsumerGroupInfoStream(String),
44    LeaseJobsBatch {
45        queue: String,
46        worker_id: String,
47        lease_seconds: i64,
48        batch_size: i64,
49    },
50    LeaseJobsBatchWithOrdering {
51        queue: String,
52        worker_id: String,
53        lease_seconds: i64,
54        batch_size: i64,
55        ordering: crate::model::QueueOrdering,
56    },
57    ReapExpiredLocks,
58    StartAttemptsBatch {
59        job_ids: Vec<Uuid>,
60        worker_id: String,
61    },
62    MarkSucceeded {
63        job_id: Uuid,
64        attempt_id: Uuid,
65        worker_id: String,
66        latency_ms: i32,
67    },
68    MarkSucceededBatch {
69        dataset_id: String,
70        updates: Vec<(Uuid, Uuid, i32)>,
71        worker_id: String,
72    },
73    RescheduleForRetry {
74        job_id: Uuid,
75        attempt_id: Uuid,
76        worker_id: String,
77        latency_ms: i32,
78        next_run_at: DateTime<Utc>,
79        error_code: String,
80        error_message: String,
81        attempt_no: i32,
82    },
83    MarkDlq {
84        job_id: Uuid,
85        attempt_id: Uuid,
86        worker_id: String,
87        latency_ms: i32,
88        reason_code: String,
89        error_code: String,
90        error_message: String,
91        attempt_no: i32,
92    },
93    ArchiveSucceededOlderThan {
94        cutoff: DateTime<Utc>,
95        limit: i64,
96    },
97    DeleteHistoryForSucceededOlderThan {
98        cutoff: DateTime<Utc>,
99        limit: i64,
100    },
101    PerformMaintenance,
102    ExtendLease {
103        job_id: Uuid,
104        worker_id: String,
105        lease_seconds: i64,
106    },
107    CancelJob {
108        job_id: Uuid,
109        worker_id: Option<String>,
110    },
111    GetJob(Uuid),
112    ListJobs {
113        queue: Option<String>,
114        status: Option<String>,
115        limit: i64,
116    },
117    ReplayJob {
118        job_id: Uuid,
119        override_queue: Option<String>,
120        override_run_at: Option<DateTime<Utc>>,
121    },
122}
123
124/// Recording mock storage backend wrapper for assertion-driven integration testing.
125#[derive(Clone)]
126pub struct MockBackend {
127    inner: Arc<dyn StorageBackend>,
128    calls: Arc<Mutex<Vec<CallRecord>>>,
129}
130
131impl MockBackend {
132    /// Creates a new `MockBackend` wrapping a target inner [`StorageBackend`].
133    pub fn new(inner: Arc<dyn StorageBackend>) -> Self {
134        Self {
135            inner,
136            calls: Arc::new(Mutex::new(Vec::new())),
137        }
138    }
139
140    /// Creates a `MockBackend` backed by an in-memory storage engine ([`MemoryBackend`]).
141    pub fn with_memory() -> Self {
142        Self::new(Arc::new(MemoryBackend::new()))
143    }
144
145    /// Returns a copy of all recorded backend call invocations.
146    pub fn calls(&self) -> Vec<CallRecord> {
147        self.calls.lock().unwrap().clone()
148    }
149
150    /// Clears recorded call history.
151    pub fn clear_calls(&self) {
152        self.calls.lock().unwrap().clear();
153    }
154
155    /// Asserts that a job of `job_type` was enqueued through this backend.
156    pub fn assert_enqueued_job_type(&self, job_type: &str) {
157        let calls = self.calls();
158        let found = calls.iter().any(|c| match c {
159            CallRecord::Enqueue(job) => job.job_type == job_type,
160            _ => false,
161        });
162        assert!(
163            found,
164            "Expected job of type '{job_type}' to be enqueued in MockBackend, but calls were: {calls:?}"
165        );
166    }
167
168    /// Asserts that a job was completed successfully via `mark_succeeded`.
169    pub fn assert_marked_succeeded(&self, target_job_id: Uuid) {
170        let calls = self.calls();
171        let found = calls.iter().any(|c| match c {
172            CallRecord::MarkSucceeded { job_id, .. } => *job_id == target_job_id,
173            CallRecord::MarkSucceededBatch { updates, .. } => {
174                updates.iter().any(|(jid, _, _)| *jid == target_job_id)
175            }
176            _ => false,
177        });
178        assert!(
179            found,
180            "Expected job {target_job_id} to be marked succeeded in MockBackend"
181        );
182    }
183
184    /// Asserts that a job was moved to Dead-Letter Queue (DLQ) via `mark_dlq`.
185    pub fn assert_marked_dlq(&self, target_job_id: Uuid) {
186        let calls = self.calls();
187        let found = calls.iter().any(|c| match c {
188            CallRecord::MarkDlq { job_id, .. } => *job_id == target_job_id,
189            _ => false,
190        });
191        assert!(
192            found,
193            "Expected job {target_job_id} to be marked DLQ in MockBackend"
194        );
195    }
196}
197
198#[async_trait]
199impl StorageBackend for MockBackend {
200    fn capabilities(&self) -> crate::model::BackendCapabilities {
201        self.inner.capabilities()
202    }
203
204    fn as_stream(&self) -> Option<&dyn StreamBackend> {
205        Some(self)
206    }
207
208    fn as_observability(&self) -> Option<&dyn ObservabilityBackend> {
209        Some(self)
210    }
211
212    async fn run_migrations(&self) -> anyhow::Result<()> {
213        self.calls.lock().unwrap().push(CallRecord::RunMigrations);
214        self.inner.run_migrations().await
215    }
216
217    async fn health_check(&self) -> anyhow::Result<()> {
218        self.calls.lock().unwrap().push(CallRecord::HealthCheck);
219        self.inner.health_check().await
220    }
221
222    async fn enqueue(&self, job: NewJob) -> anyhow::Result<Uuid> {
223        self.calls
224            .lock()
225            .unwrap()
226            .push(CallRecord::Enqueue(job.clone()));
227        self.inner.enqueue(job).await
228    }
229
230    async fn subscribe(&self, queue: &str) -> anyhow::Result<NotificationStream> {
231        self.calls
232            .lock()
233            .unwrap()
234            .push(CallRecord::Subscribe(queue.to_string()));
235        self.inner.subscribe(queue).await
236    }
237
238    async fn lease_jobs_batch(
239        &self,
240        queue: &str,
241        worker_id: &str,
242        lease_seconds: i64,
243        batch_size: i64,
244    ) -> anyhow::Result<Vec<Job>> {
245        self.calls.lock().unwrap().push(CallRecord::LeaseJobsBatch {
246            queue: queue.to_string(),
247            worker_id: worker_id.to_string(),
248            lease_seconds,
249            batch_size,
250        });
251        self.inner
252            .lease_jobs_batch(queue, worker_id, lease_seconds, batch_size)
253            .await
254    }
255
256    async fn lease_jobs_batch_with_ordering(
257        &self,
258        queue: &str,
259        worker_id: &str,
260        lease_seconds: i64,
261        batch_size: i64,
262        ordering: crate::model::QueueOrdering,
263    ) -> anyhow::Result<Vec<Job>> {
264        self.calls
265            .lock()
266            .unwrap()
267            .push(CallRecord::LeaseJobsBatchWithOrdering {
268                queue: queue.to_string(),
269                worker_id: worker_id.to_string(),
270                lease_seconds,
271                batch_size,
272                ordering,
273            });
274        self.inner
275            .lease_jobs_batch_with_ordering(queue, worker_id, lease_seconds, batch_size, ordering)
276            .await
277    }
278
279    async fn reap_expired_locks(&self) -> anyhow::Result<u64> {
280        self.calls
281            .lock()
282            .unwrap()
283            .push(CallRecord::ReapExpiredLocks);
284        self.inner.reap_expired_locks().await
285    }
286
287    async fn start_attempts_batch(
288        &self,
289        dataset_ids: &[String],
290        job_ids: &[Uuid],
291        worker_id: &str,
292    ) -> anyhow::Result<Vec<(Uuid, Uuid, i32)>> {
293        self.calls
294            .lock()
295            .unwrap()
296            .push(CallRecord::StartAttemptsBatch {
297                job_ids: job_ids.to_vec(),
298                worker_id: worker_id.to_string(),
299            });
300        self.inner
301            .start_attempts_batch(dataset_ids, job_ids, worker_id)
302            .await
303    }
304
305    async fn mark_succeeded(
306        &self,
307        job_id: Uuid,
308        attempt_id: Uuid,
309        worker_id: &str,
310        latency_ms: i32,
311    ) -> anyhow::Result<()> {
312        self.calls.lock().unwrap().push(CallRecord::MarkSucceeded {
313            job_id,
314            attempt_id,
315            worker_id: worker_id.to_string(),
316            latency_ms,
317        });
318        self.inner
319            .mark_succeeded(job_id, attempt_id, worker_id, latency_ms)
320            .await
321    }
322
323    async fn mark_succeeded_batch(
324        &self,
325        dataset_id: &str,
326        updates: &[(Uuid, Uuid, i32)],
327        worker_id: &str,
328    ) -> anyhow::Result<()> {
329        self.calls
330            .lock()
331            .unwrap()
332            .push(CallRecord::MarkSucceededBatch {
333                dataset_id: dataset_id.to_string(),
334                updates: updates.to_vec(),
335                worker_id: worker_id.to_string(),
336            });
337        self.inner
338            .mark_succeeded_batch(dataset_id, updates, worker_id)
339            .await
340    }
341
342    #[allow(clippy::too_many_arguments)]
343    async fn reschedule_for_retry(
344        &self,
345        job_id: Uuid,
346        attempt_id: Uuid,
347        worker_id: &str,
348        latency_ms: i32,
349        next_run_at: DateTime<Utc>,
350        error_code: &str,
351        error_message: &str,
352        attempt_no: i32,
353    ) -> anyhow::Result<()> {
354        self.calls
355            .lock()
356            .unwrap()
357            .push(CallRecord::RescheduleForRetry {
358                job_id,
359                attempt_id,
360                worker_id: worker_id.to_string(),
361                latency_ms,
362                next_run_at,
363                error_code: error_code.to_string(),
364                error_message: error_message.to_string(),
365                attempt_no,
366            });
367        self.inner
368            .reschedule_for_retry(
369                job_id,
370                attempt_id,
371                worker_id,
372                latency_ms,
373                next_run_at,
374                error_code,
375                error_message,
376                attempt_no,
377            )
378            .await
379    }
380
381    #[allow(clippy::too_many_arguments)]
382    async fn mark_dlq(
383        &self,
384        job_id: Uuid,
385        attempt_id: Uuid,
386        worker_id: &str,
387        latency_ms: i32,
388        reason_code: &str,
389        error_code: &str,
390        error_message: &str,
391        attempt_no: i32,
392    ) -> anyhow::Result<()> {
393        self.calls.lock().unwrap().push(CallRecord::MarkDlq {
394            job_id,
395            attempt_id,
396            worker_id: worker_id.to_string(),
397            latency_ms,
398            reason_code: reason_code.to_string(),
399            error_code: error_code.to_string(),
400            error_message: error_message.to_string(),
401            attempt_no,
402        });
403        self.inner
404            .mark_dlq(
405                job_id,
406                attempt_id,
407                worker_id,
408                latency_ms,
409                reason_code,
410                error_code,
411                error_message,
412                attempt_no,
413            )
414            .await
415    }
416
417    async fn archive_succeeded_older_than(
418        &self,
419        cutoff: DateTime<Utc>,
420        limit: i64,
421    ) -> anyhow::Result<u64> {
422        self.calls
423            .lock()
424            .unwrap()
425            .push(CallRecord::ArchiveSucceededOlderThan { cutoff, limit });
426        self.inner.archive_succeeded_older_than(cutoff, limit).await
427    }
428
429    async fn delete_history_for_succeeded_older_than(
430        &self,
431        cutoff: DateTime<Utc>,
432        limit: i64,
433    ) -> anyhow::Result<(u64, u64)> {
434        self.calls
435            .lock()
436            .unwrap()
437            .push(CallRecord::DeleteHistoryForSucceededOlderThan { cutoff, limit });
438        self.inner
439            .delete_history_for_succeeded_older_than(cutoff, limit)
440            .await
441    }
442
443    async fn perform_maintenance(&self) -> anyhow::Result<()> {
444        self.calls
445            .lock()
446            .unwrap()
447            .push(CallRecord::PerformMaintenance);
448        self.inner.perform_maintenance().await
449    }
450
451    async fn extend_lease(
452        &self,
453        job_id: Uuid,
454        worker_id: &str,
455        lease_seconds: i64,
456    ) -> anyhow::Result<bool> {
457        self.calls.lock().unwrap().push(CallRecord::ExtendLease {
458            job_id,
459            worker_id: worker_id.to_string(),
460            lease_seconds,
461        });
462        self.inner
463            .extend_lease(job_id, worker_id, lease_seconds)
464            .await
465    }
466
467    async fn cancel_job(&self, job_id: Uuid, worker_id: Option<&str>) -> anyhow::Result<()> {
468        self.calls.lock().unwrap().push(CallRecord::CancelJob {
469            job_id,
470            worker_id: worker_id.map(|id| id.to_string()),
471        });
472        self.inner.cancel_job(job_id, worker_id).await
473    }
474
475    async fn get_job(&self, job_id: Uuid) -> anyhow::Result<Option<Job>> {
476        self.calls.lock().unwrap().push(CallRecord::GetJob(job_id));
477        self.inner.get_job(job_id).await
478    }
479
480    async fn list_jobs(
481        &self,
482        queue: Option<&str>,
483        status: Option<&str>,
484        limit: i64,
485        cursor_created_at: Option<DateTime<Utc>>,
486        cursor_id: Option<Uuid>,
487    ) -> anyhow::Result<Vec<JobListItem>> {
488        self.calls.lock().unwrap().push(CallRecord::ListJobs {
489            queue: queue.map(|s| s.to_string()),
490            status: status.map(|s| s.to_string()),
491            limit,
492        });
493        self.inner
494            .list_jobs(queue, status, limit, cursor_created_at, cursor_id)
495            .await
496    }
497
498    async fn replay_job(
499        &self,
500        job_id: Uuid,
501        override_queue: Option<&str>,
502        override_run_at: Option<DateTime<Utc>>,
503    ) -> anyhow::Result<Uuid> {
504        self.calls.lock().unwrap().push(CallRecord::ReplayJob {
505            job_id,
506            override_queue: override_queue.map(|s| s.to_string()),
507            override_run_at,
508        });
509        self.inner
510            .replay_job(job_id, override_queue, override_run_at)
511            .await
512    }
513}
514
515#[async_trait]
516impl ObservabilityBackend for MockBackend {
517    async fn explain_job(&self, job_id: Uuid) -> anyhow::Result<Option<JobExplanation>> {
518        if let Some(observability) = self.inner.as_observability() {
519            observability.explain_job(job_id).await
520        } else {
521            Ok(None)
522        }
523    }
524
525    async fn queue_metrics(&self, queue: Option<&str>) -> anyhow::Result<Vec<QueueMetrics>> {
526        if let Some(observability) = self.inner.as_observability() {
527            observability.queue_metrics(queue).await
528        } else {
529            Ok(Vec::new())
530        }
531    }
532}
533
534#[async_trait]
535impl StreamBackend for MockBackend {
536    async fn publish(&self, stream: &str, event: NewEvent) -> anyhow::Result<i64> {
537        self.calls.lock().unwrap().push(CallRecord::PublishStream {
538            stream: stream.to_string(),
539            event: event.clone(),
540        });
541        if let Some(sb) = self.inner.as_stream() {
542            sb.publish(stream, event).await
543        } else {
544            Ok(1)
545        }
546    }
547
548    async fn subscribe_stream(
549        &self,
550        stream: &str,
551        consumer_group: &str,
552        last_seq: Option<i64>,
553    ) -> anyhow::Result<NotificationStream> {
554        self.calls
555            .lock()
556            .unwrap()
557            .push(CallRecord::SubscribeStream {
558                stream: stream.to_string(),
559                consumer_group: consumer_group.to_string(),
560                last_seq,
561            });
562        if let Some(sb) = self.inner.as_stream() {
563            sb.subscribe_stream(stream, consumer_group, last_seq).await
564        } else {
565            anyhow::bail!("inner backend does not support StreamBackend")
566        }
567    }
568
569    async fn ack(&self, stream: &str, consumer_group: &str, seq: i64) -> anyhow::Result<()> {
570        self.calls.lock().unwrap().push(CallRecord::AckStream {
571            stream: stream.to_string(),
572            consumer_group: consumer_group.to_string(),
573            seq,
574        });
575        if let Some(sb) = self.inner.as_stream() {
576            sb.ack(stream, consumer_group, seq).await
577        } else {
578            Ok(())
579        }
580    }
581
582    async fn read_events(
583        &self,
584        stream: &str,
585        after_seq: i64,
586        limit: i64,
587    ) -> anyhow::Result<Vec<Event>> {
588        self.calls
589            .lock()
590            .unwrap()
591            .push(CallRecord::ReadEventsStream {
592                stream: stream.to_string(),
593                after_seq,
594                limit,
595            });
596        if let Some(sb) = self.inner.as_stream() {
597            sb.read_events(stream, after_seq, limit).await
598        } else {
599            Ok(Vec::new())
600        }
601    }
602
603    async fn prune_events(&self, stream: &str, through_seq: i64) -> anyhow::Result<u64> {
604        self.calls
605            .lock()
606            .unwrap()
607            .push(CallRecord::PruneEventsStream {
608                stream: stream.to_string(),
609                through_seq,
610            });
611        if let Some(sb) = self.inner.as_stream() {
612            sb.prune_events(stream, through_seq).await
613        } else {
614            Ok(0)
615        }
616    }
617
618    async fn consumer_group_info(&self, stream: &str) -> anyhow::Result<Vec<ConsumerGroupStatus>> {
619        self.calls
620            .lock()
621            .unwrap()
622            .push(CallRecord::ConsumerGroupInfoStream(stream.to_string()));
623        if let Some(sb) = self.inner.as_stream() {
624            sb.consumer_group_info(stream).await
625        } else {
626            Ok(Vec::new())
627        }
628    }
629}
azums_core/backend/mock.rs

azums_core/backend/
mock.rs
