azums_core/backend/
mock.rs

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

azums_core/backend/
mock.rs
