apalis_core/worker/
test_worker.rs

1//! Provides a worker that allows testing and debugging
2//!
3//! This module enables comprehensive testing of task services and backends by simulating
4//! a real worker’s lifecycle. It allows developers to push jobs to a backend, run them
5//! through a test worker, and capture the results for assertion without needing full runtime orchestration.
6//!
7//! # Features
8//! - Pluggable with any backend implementing [`Backend`].
9//! - Supports service functions created using [`Service`].
10//! - Captures task output through [`TestEmitService`] for validation.
11//!
12//! # Example
13//! ```rust
14//! # use apalis_core::{backend::memory::MemoryStorage, worker::test_worker::TestWorker};
15//! # use apalis_core::error::BoxDynError;
16//! async fn is_even(req: usize) -> Result<(), BoxDynError> {
17//!     if req % 2 == 0 {
18//!         Ok(())
19//!     } else {
20//!         Err("Not an even number".into())
21//!     }
22//! }
23//!
24//! #[tokio::test]
25//! async fn test_accepts_even() {
26//!     let backend = MemoryStorage::new();
27//!     backend.push(42usize).await.unwrap();
28//!     let mut worker = TestWorker::new(backend, is_even);
29//!
30//!     let (_task_id, resp) = worker.execute_next().await.unwrap().unwrap();
31//!     assert_eq!(resp, Ok("()".to_string()));
32//! }
33//! ```
34//!
35//! This module is intended for use in tests and local development.
36
37use crate::backend::Backend;
38use crate::error::BoxDynError;
39use crate::task::Task;
40use crate::task::task_id::{RandomId, TaskId};
41use crate::worker::builder::{IntoWorkerService, WorkerBuilder};
42use crate::worker::{Event, ReadinessService, TrackerService, WorkerError};
43use futures_channel::mpsc::{self, channel};
44use futures_core::future::BoxFuture;
45use futures_core::stream::BoxStream;
46use futures_util::{SinkExt, StreamExt};
47use std::fmt::{self, Debug};
48use std::future::Future;
49use std::marker::PhantomData;
50use std::task::{Context, Poll};
51use tower_layer::Layer;
52use tower_service::Service;
53
54type TestStream<IdType, Res> =
55    BoxStream<'static, Result<(TaskId<IdType>, Result<Res, BoxDynError>), WorkerError>>;
56/// A test worker to allow you to test services.
57/// Important for testing backends and tasks
58/// # Example
59/// ```
60/// mod tests {
61///    use apalis_core::{error::BoxDynError, backend::memory::MemoryStorage};
62///
63///    use super::*;
64///
65///    async fn is_even(req: usize) -> Result<(), BoxDynError> {
66///        if req % 2 == 0 {
67///            Ok(())
68///        } else {
69///            Err("Not an even number".into())
70///        }
71///    }
72///
73///    #[tokio::test]
74///    async fn test_accepts_even() {
75///        let mut backend = MemoryStorage::new();
76///        backend.push(42usize).await.unwrap();
77///        let mut worker = TestWorker::new(backend, is_even);
78///        let (_task_id, resp) = worker.execute_next().await.unwrap().unwrap();
79///        assert_eq!(resp, Ok("()".to_string()));
80///    }
81///}
82/// ````
83pub struct TestWorker<B, S, Res, IdType = RandomId> {
84    stream: TestStream<IdType, Res>,
85    backend: PhantomData<B>,
86    service: PhantomData<(S, Res)>,
87}
88
89impl<B, S, Res, IdType> fmt::Debug for TestWorker<B, S, Res, IdType> {
90    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91        f.debug_struct("TestWorker")
92            .field("stream", &"BoxStream<...>") // can't really debug streams
93            .field("backend", &std::any::type_name::<B>())
94            .field("service", &std::any::type_name::<(S, Res)>())
95            .field("id_type", &std::any::type_name::<IdType>())
96            .finish()
97    }
98}
99
100/// Utility for executing the next item in the queue
101pub trait ExecuteNext<Args, Ctx> {
102    /// The expected result from the provided service
103    type Result;
104    /// Allows the test worker to step to the next task
105    /// No polling is done in between calls
106    fn execute_next(&mut self) -> impl Future<Output = Self::Result> + Send;
107}
108
109impl<B, S, Args, Ctx, Res, IdType> ExecuteNext<Args, Ctx> for TestWorker<B, S, Res, IdType>
110where
111    S: Service<Task<Args, Ctx, IdType>, Response = Res> + Send + 'static,
112    B: Send,
113    Res: Send,
114{
115    type Result = Option<Result<(TaskId<IdType>, Result<Res, BoxDynError>), WorkerError>>;
116    async fn execute_next(&mut self) -> Self::Result {
117        self.stream.next().await
118    }
119}
120
121impl<B, S, Res> TestWorker<B, S, Res, ()> {
122    /// Create a new test worker
123    pub fn new<Args, Ctx, W>(backend: B, factory: W) -> TestWorker<W::Backend, S, Res, B::IdType>
124    where
125        W: IntoWorkerService<B, S, Args, Ctx>,
126        W::Backend: Backend<
127                Args = B::Args,
128                Context = B::Context,
129                IdType = B::IdType,
130                Error = B::Error,
131                Stream = B::Stream,
132                Beat = B::Beat,
133                Layer = B::Layer,
134            > + 'static,
135        B: Backend<Args = Args, Context = Ctx> + 'static,
136        S: Service<Task<Args, Ctx, B::IdType>, Response = Res> + Send + 'static,
137        B::Stream: Unpin + Send + 'static,
138        B::Beat: Unpin + Send + 'static,
139        Args: Send + 'static,
140        Ctx: Send + 'static,
141        B::Error: Into<BoxDynError> + Send + 'static,
142        B::Layer: Layer<ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>>,
143        S::Future: Send,
144        S::Error: Into<BoxDynError> + Send + Sync,
145        S::Response: Clone + Send,
146        Res: 'static,
147        <<B as Backend>::Layer as Layer<
148            ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>,
149        >>::Service: Service<Task<Args, Ctx, B::IdType>>,
150        <<<B as Backend>::Layer as Layer<
151            ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>,
152        >>::Service as Service<Task<Args, Ctx, B::IdType>>>::Error: Into<BoxDynError> + Sync + Send,
153        <<<B as Backend>::Layer as Layer<
154            ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>,
155        >>::Service as Service<Task<Args, Ctx, B::IdType>>>::Future: Send,
156        <<B as Backend>::Layer as Layer<
157            ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>,
158        >>::Service: std::marker::Send + 'static,
159        B::IdType: Send + Clone + 'static,
160    {
161        let worker_service = factory.into_service(backend);
162        TestWorker::<W::Backend, S, Res, _>::new_with_svc(
163            worker_service.backend,
164            worker_service.service,
165        )
166    }
167}
168
169impl<B, S, Res> TestWorker<B, S, Res, ()> {
170    /// Create a new test worker with a service
171    pub fn new_with_svc<Args, Ctx>(backend: B, service: S) -> TestWorker<B, S, Res, B::IdType>
172    where
173        B: Backend<Args = Args, Context = Ctx> + 'static,
174        S: Service<Task<Args, Ctx, B::IdType>, Response = Res> + Send + 'static,
175        B::Stream: Unpin + Send + 'static,
176        B::Beat: Unpin + Send + 'static,
177        Args: Send + 'static,
178        Ctx: Send + 'static,
179        B::Error: Into<BoxDynError> + Send + 'static,
180        B::Layer: Layer<ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>>,
181        S::Future: Send,
182        S::Error: Into<BoxDynError> + Send + Sync,
183        S::Response: Clone + Send,
184        Res: 'static,
185        <<B as Backend>::Layer as Layer<
186            ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>,
187        >>::Service: Service<Task<Args, Ctx, B::IdType>>,
188        <<<B as Backend>::Layer as Layer<
189            ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>,
190        >>::Service as Service<Task<Args, Ctx, B::IdType>>>::Error: Into<BoxDynError> + Sync + Send,
191        <<<B as Backend>::Layer as Layer<
192            ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>,
193        >>::Service as Service<Task<Args, Ctx, B::IdType>>>::Future: Send,
194        <<B as Backend>::Layer as Layer<
195            ReadinessService<TrackerService<TestEmitService<S, Res, B::IdType>>>,
196        >>::Service: std::marker::Send + 'static,
197        B::IdType: Send + Clone + 'static,
198    {
199        enum Item<R, IdType> {
200            Ev(Result<Event, WorkerError>),
201            Res((TaskId<IdType>, Result<R, BoxDynError>)),
202        }
203        let (tx, rx) = channel(1);
204        let sender = tx.clone();
205        let service: TestEmitService<S, Res, B::IdType> = TestEmitService {
206            service,
207            tx: tx.clone(),
208        };
209        let stream = WorkerBuilder::new("test-worker")
210            .backend(backend)
211            .build(service)
212            .stream()
213            .map(|r| Item::Ev(r));
214        let task_stream = rx.map(|s| Item::Res(s));
215        let stream = futures_util::stream::select(task_stream, stream)
216            .filter_map(move |s| {
217                let mut tx = sender.clone();
218                async move {
219                    match s {
220                        Item::Ev(Err(e)) => {
221                            tx.close().await.unwrap();
222                            Some(Err(e))
223                        }
224                        Item::Ev(_) => None,
225                        Item::Res(r) => Some(Ok(r)),
226                    }
227                }
228            })
229            .boxed();
230        TestWorker {
231            stream,
232            service: PhantomData,
233            backend: PhantomData,
234        }
235    }
236}
237
238/// A generic service that emits the result of a test
239#[derive(Debug, Clone)]
240pub struct TestEmitService<S, Response, IdType> {
241    tx: mpsc::Sender<(TaskId<IdType>, Result<Response, BoxDynError>)>,
242    service: S,
243}
244
245impl<S, Args, Ctx, Res, IdType> Service<Task<Args, Ctx, IdType>> for TestEmitService<S, Res, IdType>
246where
247    S: Service<Task<Args, Ctx, IdType>, Response = Res> + Send + 'static,
248    S::Future: Send + 'static,
249    Args: Send + 'static,
250    Ctx: Send + 'static,
251    S::Response: Send + 'static,
252    S::Error: Into<BoxDynError> + Send,
253    IdType: Send + 'static + Clone,
254{
255    type Response = ();
256    type Error = String;
257    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
258
259    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
260        self.service
261            .poll_ready(cx)
262            .map_err(|e| e.into().to_string())
263    }
264
265    fn call(&mut self, task: Task<Args, Ctx, IdType>) -> Self::Future {
266        let task_id = task.parts.task_id.clone().unwrap();
267        let mut tx = Clone::clone(&self.tx);
268        let fut = self.service.call(task);
269        Box::pin(async move {
270            let res = fut.await;
271            match res {
272                Ok(res) => {
273                    tx.send((task_id, Ok(res))).await.unwrap();
274                    Ok(())
275                }
276                Err(err) => {
277                    let e = err.into();
278                    let e_str = e.to_string();
279                    tx.send((task_id, Err(e))).await.unwrap();
280                    Err(e_str)
281                }
282            }
283        })
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use crate::{
290        backend::{TaskSink, memory::MemoryStorage},
291        error::BoxDynError,
292        task_fn::task_fn,
293        worker::{
294            WorkerContext,
295            test_worker::{ExecuteNext, TestWorker},
296        },
297    };
298    use std::time::Duration;
299
300    #[tokio::test]
301    async fn basic_worker() {
302        let mut backend = MemoryStorage::new();
303
304        for i in 0..=10 {
305            backend.push(i).await.unwrap();
306        }
307
308        let service = task_fn(|req: u32, w: WorkerContext| async move {
309            if req == 10 {
310                w.stop()?;
311            }
312            tokio::time::sleep(Duration::from_secs(1)).await;
313            Ok::<_, BoxDynError>(req)
314        });
315        let mut worker = TestWorker::new(backend, service);
316        while let Some(Ok((_, ret))) = worker.execute_next().await {
317            ret.unwrap();
318        }
319        println!("Worker run successfully");
320    }
321}
apalis_core/worker/test_worker.rs

apalis_core/worker/
test_worker.rs
