forge_core/testing/context/
job.rs

1//! Test context for job functions.
2
3use std::collections::HashMap;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::sync::{Arc, RwLock};
6
7use sqlx::PgPool;
8use uuid::Uuid;
9
10use super::super::mock_http::{MockHttp, MockRequest, MockResponse};
11use super::build_test_auth;
12use crate::Result;
13use crate::env::{EnvAccess, EnvProvider, MockEnvProvider};
14use crate::function::AuthContext;
15
16/// Progress update recorded during testing.
17#[derive(Debug, Clone)]
18pub struct TestProgressUpdate {
19    /// Progress percentage (0-100).
20    pub percent: u8,
21    /// Progress message.
22    pub message: String,
23}
24
25/// Test context for job functions.
26///
27/// Provides an isolated testing environment for jobs with progress tracking,
28/// retry simulation, cancellation testing, and HTTP mocking.
29///
30/// # Example
31///
32/// ```ignore
33/// let ctx = TestJobContext::builder("export_users")
34///     .with_job_id(Uuid::new_v4())
35///     .build();
36///
37/// // Simulate progress
38/// ctx.progress(50, "Halfway there")?;
39///
40/// // Verify progress was recorded
41/// assert_eq!(ctx.progress_updates().len(), 1);
42///
43/// // Test cancellation
44/// ctx.request_cancellation();
45/// assert!(ctx.is_cancel_requested().unwrap());
46/// ```
47pub struct TestJobContext {
48    /// Job ID.
49    pub job_id: Uuid,
50    /// Job type name.
51    pub job_type: String,
52    /// Current attempt number (1-based).
53    pub attempt: u32,
54    /// Maximum attempts allowed.
55    pub max_attempts: u32,
56    /// Authentication context.
57    pub auth: AuthContext,
58    /// Optional database pool.
59    pool: Option<PgPool>,
60    /// Mock HTTP client.
61    http: Arc<MockHttp>,
62    /// Progress updates recorded during execution.
63    progress_updates: Arc<RwLock<Vec<TestProgressUpdate>>>,
64    /// Mock environment provider.
65    env_provider: Arc<MockEnvProvider>,
66    /// Persisted saved data (in-memory).
67    saved_data: Arc<RwLock<serde_json::Value>>,
68    /// Whether cancellation has been requested.
69    cancel_requested: Arc<AtomicBool>,
70}
71
72impl TestJobContext {
73    /// Create a new builder.
74    pub fn builder(job_type: impl Into<String>) -> TestJobContextBuilder {
75        TestJobContextBuilder::new(job_type)
76    }
77
78    /// Get the database pool (if available).
79    pub fn db(&self) -> Option<&PgPool> {
80        self.pool.as_ref()
81    }
82
83    /// Get the mock HTTP client.
84    pub fn http(&self) -> &MockHttp {
85        &self.http
86    }
87
88    /// Report job progress.
89    pub fn progress(&self, percent: u8, message: impl Into<String>) -> Result<()> {
90        let update = TestProgressUpdate {
91            percent: percent.min(100),
92            message: message.into(),
93        };
94        self.progress_updates.write().unwrap().push(update);
95        Ok(())
96    }
97
98    /// Get all progress updates.
99    pub fn progress_updates(&self) -> Vec<TestProgressUpdate> {
100        self.progress_updates.read().unwrap().clone()
101    }
102
103    /// Get all saved job data.
104    pub fn saved(&self) -> serde_json::Value {
105        self.saved_data.read().unwrap().clone()
106    }
107
108    /// Replace all saved job data.
109    pub fn set_saved(&self, data: serde_json::Value) -> Result<()> {
110        let mut guard = self.saved_data.write().unwrap();
111        *guard = data;
112        Ok(())
113    }
114
115    /// Save a key-value pair to job data.
116    pub fn save(&self, key: &str, value: serde_json::Value) -> Result<()> {
117        let mut guard = self.saved_data.write().unwrap();
118        if let Some(map) = guard.as_object_mut() {
119            map.insert(key.to_string(), value);
120        } else {
121            let mut map = serde_json::Map::new();
122            map.insert(key.to_string(), value);
123            *guard = serde_json::Value::Object(map);
124        }
125        Ok(())
126    }
127
128    /// Check if this is a retry attempt.
129    pub fn is_retry(&self) -> bool {
130        self.attempt > 1
131    }
132
133    /// Check if this is the last attempt.
134    pub fn is_last_attempt(&self) -> bool {
135        self.attempt >= self.max_attempts
136    }
137
138    /// Simulate heartbeat (no-op in tests, but records the intent).
139    pub async fn heartbeat(&self) -> Result<()> {
140        Ok(())
141    }
142
143    /// Check if cancellation has been requested.
144    pub fn is_cancel_requested(&self) -> Result<bool> {
145        Ok(self.cancel_requested.load(Ordering::SeqCst))
146    }
147
148    /// Return an error if cancellation has been requested.
149    ///
150    /// Use this in job handlers to check for cancellation and exit early.
151    pub fn check_cancelled(&self) -> Result<()> {
152        if self.cancel_requested.load(Ordering::SeqCst) {
153            Err(crate::ForgeError::JobCancelled(
154                "Job cancellation requested".to_string(),
155            ))
156        } else {
157            Ok(())
158        }
159    }
160
161    /// Request cancellation (for testing cancellation flows).
162    ///
163    /// After calling this, `is_cancel_requested()` returns `true` and
164    /// `check_cancelled()` returns an error.
165    pub fn request_cancellation(&self) {
166        self.cancel_requested.store(true, Ordering::SeqCst);
167    }
168
169    /// Get the mock env provider for verification.
170    pub fn env_mock(&self) -> &MockEnvProvider {
171        &self.env_provider
172    }
173}
174
175impl EnvAccess for TestJobContext {
176    fn env_provider(&self) -> &dyn EnvProvider {
177        self.env_provider.as_ref()
178    }
179}
180
181/// Builder for TestJobContext.
182pub struct TestJobContextBuilder {
183    job_id: Option<Uuid>,
184    job_type: String,
185    attempt: u32,
186    max_attempts: u32,
187    user_id: Option<Uuid>,
188    roles: Vec<String>,
189    claims: HashMap<String, serde_json::Value>,
190    pool: Option<PgPool>,
191    http: MockHttp,
192    env_vars: HashMap<String, String>,
193    cancel_requested: bool,
194}
195
196impl TestJobContextBuilder {
197    /// Create a new builder with job type.
198    pub fn new(job_type: impl Into<String>) -> Self {
199        Self {
200            job_id: None,
201            job_type: job_type.into(),
202            attempt: 1,
203            max_attempts: 1,
204            user_id: None,
205            roles: Vec::new(),
206            claims: HashMap::new(),
207            pool: None,
208            http: MockHttp::new(),
209            env_vars: HashMap::new(),
210            cancel_requested: false,
211        }
212    }
213
214    /// Set a specific job ID.
215    pub fn with_job_id(mut self, id: Uuid) -> Self {
216        self.job_id = Some(id);
217        self
218    }
219
220    /// Set as a retry (attempt > 1).
221    pub fn as_retry(mut self, attempt: u32) -> Self {
222        self.attempt = attempt.max(1);
223        self
224    }
225
226    /// Set the maximum attempts.
227    pub fn with_max_attempts(mut self, max: u32) -> Self {
228        self.max_attempts = max.max(1);
229        self
230    }
231
232    /// Set as the last attempt.
233    pub fn as_last_attempt(mut self) -> Self {
234        self.attempt = 3;
235        self.max_attempts = 3;
236        self
237    }
238
239    /// Set the authenticated user with a UUID.
240    pub fn as_user(mut self, id: Uuid) -> Self {
241        self.user_id = Some(id);
242        self
243    }
244
245    /// For non-UUID auth providers (Firebase, Clerk, etc.).
246    pub fn as_subject(mut self, subject: impl Into<String>) -> Self {
247        self.claims
248            .insert("sub".to_string(), serde_json::json!(subject.into()));
249        self
250    }
251
252    /// Add a role.
253    pub fn with_role(mut self, role: impl Into<String>) -> Self {
254        self.roles.push(role.into());
255        self
256    }
257
258    /// Set the database pool.
259    pub fn with_pool(mut self, pool: PgPool) -> Self {
260        self.pool = Some(pool);
261        self
262    }
263
264    /// Add an HTTP mock with a custom handler.
265    pub fn mock_http<F>(self, pattern: &str, handler: F) -> Self
266    where
267        F: Fn(&MockRequest) -> MockResponse + Send + Sync + 'static,
268    {
269        self.http.add_mock_sync(pattern, handler);
270        self
271    }
272
273    /// Add an HTTP mock that returns a JSON response.
274    pub fn mock_http_json<T: serde::Serialize>(self, pattern: &str, response: T) -> Self {
275        let json = serde_json::to_value(response).unwrap_or(serde_json::Value::Null);
276        self.mock_http(pattern, move |_| MockResponse::json(json.clone()))
277    }
278
279    /// Set a single environment variable.
280    pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
281        self.env_vars.insert(key.into(), value.into());
282        self
283    }
284
285    /// Set multiple environment variables.
286    pub fn with_envs(mut self, vars: HashMap<String, String>) -> Self {
287        self.env_vars.extend(vars);
288        self
289    }
290
291    /// Start with cancellation already requested.
292    ///
293    /// Use this to test how jobs handle cancellation from the start.
294    pub fn with_cancellation_requested(mut self) -> Self {
295        self.cancel_requested = true;
296        self
297    }
298
299    /// Build the test context.
300    pub fn build(self) -> TestJobContext {
301        TestJobContext {
302            job_id: self.job_id.unwrap_or_else(Uuid::new_v4),
303            job_type: self.job_type,
304            attempt: self.attempt,
305            max_attempts: self.max_attempts,
306            auth: build_test_auth(self.user_id, self.roles, self.claims),
307            pool: self.pool,
308            http: Arc::new(self.http),
309            progress_updates: Arc::new(RwLock::new(Vec::new())),
310            env_provider: Arc::new(MockEnvProvider::with_vars(self.env_vars)),
311            saved_data: Arc::new(RwLock::new(crate::job::empty_saved_data())),
312            cancel_requested: Arc::new(AtomicBool::new(self.cancel_requested)),
313        }
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320
321    #[test]
322    fn test_job_context_creation() {
323        let ctx = TestJobContext::builder("export_users").build();
324
325        assert_eq!(ctx.job_type, "export_users");
326        assert_eq!(ctx.attempt, 1);
327        assert!(!ctx.is_retry());
328        assert!(ctx.is_last_attempt()); // 1 of 1
329    }
330
331    #[test]
332    fn test_retry_detection() {
333        let ctx = TestJobContext::builder("test")
334            .as_retry(3)
335            .with_max_attempts(5)
336            .build();
337
338        assert!(ctx.is_retry());
339        assert!(!ctx.is_last_attempt());
340    }
341
342    #[test]
343    fn test_last_attempt() {
344        let ctx = TestJobContext::builder("test").as_last_attempt().build();
345
346        assert!(ctx.is_retry());
347        assert!(ctx.is_last_attempt());
348    }
349
350    #[test]
351    fn test_progress_tracking() {
352        let ctx = TestJobContext::builder("test").build();
353
354        ctx.progress(25, "Step 1 complete").unwrap();
355        ctx.progress(50, "Step 2 complete").unwrap();
356        ctx.progress(100, "Done").unwrap();
357
358        let updates = ctx.progress_updates();
359        assert_eq!(updates.len(), 3);
360        assert_eq!(updates[0].percent, 25);
361        assert_eq!(updates[2].percent, 100);
362    }
363
364    #[test]
365    fn test_save_and_saved() {
366        let ctx = TestJobContext::builder("test").build();
367        ctx.save("charge_id", serde_json::json!("ch_123")).unwrap();
368        ctx.save("amount", serde_json::json!(100)).unwrap();
369
370        let saved = ctx.saved();
371        assert_eq!(saved["charge_id"], "ch_123");
372        assert_eq!(saved["amount"], 100);
373    }
374
375    #[test]
376    fn test_cancellation_not_requested() {
377        let ctx = TestJobContext::builder("test").build();
378
379        assert!(!ctx.is_cancel_requested().unwrap());
380        assert!(ctx.check_cancelled().is_ok());
381    }
382
383    #[test]
384    fn test_cancellation_requested_at_build() {
385        let ctx = TestJobContext::builder("test")
386            .with_cancellation_requested()
387            .build();
388
389        assert!(ctx.is_cancel_requested().unwrap());
390        assert!(ctx.check_cancelled().is_err());
391    }
392
393    #[test]
394    fn test_request_cancellation_mid_test() {
395        let ctx = TestJobContext::builder("test").build();
396
397        assert!(!ctx.is_cancel_requested().unwrap());
398        ctx.request_cancellation();
399        assert!(ctx.is_cancel_requested().unwrap());
400        assert!(ctx.check_cancelled().is_err());
401    }
402}
forge_core/testing/context/job.rs

forge_core/testing/context/
job.rs
