forge_core/testing/context/
job.rs1use std::collections::HashMap;
4use std::sync::{Arc, RwLock};
5
6use sqlx::PgPool;
7use uuid::Uuid;
8
9use super::super::mock_http::{MockHttp, MockRequest, MockResponse};
10use crate::Result;
11use crate::env::{EnvAccess, EnvProvider, MockEnvProvider};
12use crate::function::AuthContext;
13
14#[derive(Debug, Clone)]
16pub struct TestProgressUpdate {
17 pub percent: u8,
19 pub message: String,
21}
22
23pub struct TestJobContext {
42 pub job_id: Uuid,
44 pub job_type: String,
46 pub attempt: u32,
48 pub max_attempts: u32,
50 pub auth: AuthContext,
52 pool: Option<PgPool>,
54 http: Arc<MockHttp>,
56 progress_updates: Arc<RwLock<Vec<TestProgressUpdate>>>,
58 env_provider: Arc<MockEnvProvider>,
60}
61
62impl TestJobContext {
63 pub fn builder(job_type: impl Into<String>) -> TestJobContextBuilder {
65 TestJobContextBuilder::new(job_type)
66 }
67
68 pub fn db(&self) -> Option<&PgPool> {
70 self.pool.as_ref()
71 }
72
73 pub fn http(&self) -> &MockHttp {
75 &self.http
76 }
77
78 pub fn progress(&self, percent: u8, message: impl Into<String>) -> Result<()> {
80 let update = TestProgressUpdate {
81 percent: percent.min(100),
82 message: message.into(),
83 };
84 self.progress_updates.write().unwrap().push(update);
85 Ok(())
86 }
87
88 pub fn progress_updates(&self) -> Vec<TestProgressUpdate> {
90 self.progress_updates.read().unwrap().clone()
91 }
92
93 pub fn is_retry(&self) -> bool {
95 self.attempt > 1
96 }
97
98 pub fn is_last_attempt(&self) -> bool {
100 self.attempt >= self.max_attempts
101 }
102
103 pub async fn heartbeat(&self) -> Result<()> {
105 Ok(())
106 }
107
108 pub fn env_mock(&self) -> &MockEnvProvider {
110 &self.env_provider
111 }
112}
113
114impl EnvAccess for TestJobContext {
115 fn env_provider(&self) -> &dyn EnvProvider {
116 self.env_provider.as_ref()
117 }
118}
119
120pub struct TestJobContextBuilder {
122 job_id: Option<Uuid>,
123 job_type: String,
124 attempt: u32,
125 max_attempts: u32,
126 user_id: Option<Uuid>,
127 roles: Vec<String>,
128 claims: HashMap<String, serde_json::Value>,
129 pool: Option<PgPool>,
130 http: MockHttp,
131 env_vars: HashMap<String, String>,
132}
133
134impl TestJobContextBuilder {
135 pub fn new(job_type: impl Into<String>) -> Self {
137 Self {
138 job_id: None,
139 job_type: job_type.into(),
140 attempt: 1,
141 max_attempts: 1,
142 user_id: None,
143 roles: Vec::new(),
144 claims: HashMap::new(),
145 pool: None,
146 http: MockHttp::new(),
147 env_vars: HashMap::new(),
148 }
149 }
150
151 pub fn with_job_id(mut self, id: Uuid) -> Self {
153 self.job_id = Some(id);
154 self
155 }
156
157 pub fn as_retry(mut self, attempt: u32) -> Self {
159 self.attempt = attempt.max(1);
160 self
161 }
162
163 pub fn with_max_attempts(mut self, max: u32) -> Self {
165 self.max_attempts = max.max(1);
166 self
167 }
168
169 pub fn as_last_attempt(mut self) -> Self {
171 self.attempt = 3;
172 self.max_attempts = 3;
173 self
174 }
175
176 pub fn as_user(mut self, id: Uuid) -> Self {
178 self.user_id = Some(id);
179 self
180 }
181
182 pub fn with_role(mut self, role: impl Into<String>) -> Self {
184 self.roles.push(role.into());
185 self
186 }
187
188 pub fn with_pool(mut self, pool: PgPool) -> Self {
190 self.pool = Some(pool);
191 self
192 }
193
194 pub fn mock_http<F>(self, pattern: &str, handler: F) -> Self
196 where
197 F: Fn(&MockRequest) -> MockResponse + Send + Sync + 'static,
198 {
199 self.http.add_mock_sync(pattern, handler);
200 self
201 }
202
203 pub fn mock_http_json<T: serde::Serialize>(self, pattern: &str, response: T) -> Self {
205 let json = serde_json::to_value(response).unwrap_or(serde_json::Value::Null);
206 self.mock_http(pattern, move |_| MockResponse::json(json.clone()))
207 }
208
209 pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
211 self.env_vars.insert(key.into(), value.into());
212 self
213 }
214
215 pub fn with_envs(mut self, vars: HashMap<String, String>) -> Self {
217 self.env_vars.extend(vars);
218 self
219 }
220
221 pub fn build(self) -> TestJobContext {
223 let auth = if let Some(user_id) = self.user_id {
224 AuthContext::authenticated(user_id, self.roles, self.claims)
225 } else {
226 AuthContext::unauthenticated()
227 };
228
229 TestJobContext {
230 job_id: self.job_id.unwrap_or_else(Uuid::new_v4),
231 job_type: self.job_type,
232 attempt: self.attempt,
233 max_attempts: self.max_attempts,
234 auth,
235 pool: self.pool,
236 http: Arc::new(self.http),
237 progress_updates: Arc::new(RwLock::new(Vec::new())),
238 env_provider: Arc::new(MockEnvProvider::with_vars(self.env_vars)),
239 }
240 }
241}
242
243#[cfg(test)]
244mod tests {
245 use super::*;
246
247 #[test]
248 fn test_job_context_creation() {
249 let ctx = TestJobContext::builder("export_users").build();
250
251 assert_eq!(ctx.job_type, "export_users");
252 assert_eq!(ctx.attempt, 1);
253 assert!(!ctx.is_retry());
254 assert!(ctx.is_last_attempt()); }
256
257 #[test]
258 fn test_retry_detection() {
259 let ctx = TestJobContext::builder("test")
260 .as_retry(3)
261 .with_max_attempts(5)
262 .build();
263
264 assert!(ctx.is_retry());
265 assert!(!ctx.is_last_attempt());
266 }
267
268 #[test]
269 fn test_last_attempt() {
270 let ctx = TestJobContext::builder("test").as_last_attempt().build();
271
272 assert!(ctx.is_retry());
273 assert!(ctx.is_last_attempt());
274 }
275
276 #[test]
277 fn test_progress_tracking() {
278 let ctx = TestJobContext::builder("test").build();
279
280 ctx.progress(25, "Step 1 complete").unwrap();
281 ctx.progress(50, "Step 2 complete").unwrap();
282 ctx.progress(100, "Done").unwrap();
283
284 let updates = ctx.progress_updates();
285 assert_eq!(updates.len(), 3);
286 assert_eq!(updates[0].percent, 25);
287 assert_eq!(updates[2].percent, 100);
288 }
289}