1use std::collections::HashMap;
4use std::time::Duration;
5
6use chrono::{DateTime, Utc};
7use uuid::Uuid;
8
9use forge_core::error::{ForgeError, Result};
10use forge_core::function::AuthContext;
11use forge_core::job::JobStatus;
12use forge_core::workflow::WorkflowStatus;
13
14use super::TestConfig;
15use super::mock::{MockHttp, MockRequest, MockResponse};
16
17pub struct TestContext {
22 pool: Option<sqlx::PgPool>,
24 mock_http: MockHttp,
26 auth: AuthContext,
28 dispatched_jobs: Vec<DispatchedJob>,
30 started_workflows: Vec<StartedWorkflow>,
32}
33
34#[derive(Debug, Clone)]
36pub struct DispatchedJob {
37 pub id: Uuid,
39 pub job_type: String,
41 pub input: serde_json::Value,
43 pub dispatched_at: DateTime<Utc>,
45 pub status: JobStatus,
47}
48
49#[derive(Debug, Clone)]
51pub struct StartedWorkflow {
52 pub run_id: Uuid,
54 pub workflow_name: String,
56 pub input: serde_json::Value,
58 pub started_at: DateTime<Utc>,
60 pub status: WorkflowStatus,
62 pub completed_steps: Vec<String>,
64}
65
66impl TestContext {
67 pub fn new_without_db() -> Self {
69 Self {
70 pool: None,
71 mock_http: MockHttp::new(),
72 auth: AuthContext::unauthenticated(),
73 dispatched_jobs: Vec::new(),
74 started_workflows: Vec::new(),
75 }
76 }
77
78 pub async fn new() -> Result<Self> {
80 let config = TestConfig::default();
81 Self::with_config(config).await
82 }
83
84 pub async fn with_config(config: TestConfig) -> Result<Self> {
86 let pool = if let Some(ref url) = config.database_url {
87 Some(
88 sqlx::postgres::PgPoolOptions::new()
89 .max_connections(config.max_connections)
90 .acquire_timeout(Duration::from_secs(30))
91 .connect(url)
92 .await
93 .map_err(|e| ForgeError::Database(e.to_string()))?,
94 )
95 } else {
96 None
97 };
98
99 Ok(Self {
100 pool,
101 mock_http: MockHttp::new(),
102 auth: AuthContext::unauthenticated(),
103 dispatched_jobs: Vec::new(),
104 started_workflows: Vec::new(),
105 })
106 }
107
108 pub fn builder() -> TestContextBuilder {
110 TestContextBuilder::new()
111 }
112
113 pub fn pool(&self) -> Option<&sqlx::PgPool> {
115 self.pool.as_ref()
116 }
117
118 pub fn auth(&self) -> &AuthContext {
120 &self.auth
121 }
122
123 pub fn user_id(&self) -> Option<Uuid> {
125 if self.auth.is_authenticated() {
126 self.auth.user_id()
127 } else {
128 None
129 }
130 }
131
132 pub fn set_user(&mut self, user_id: Uuid) {
134 self.auth = AuthContext::authenticated(user_id, vec![], HashMap::new());
135 }
136
137 pub fn mock_http(&self) -> &MockHttp {
139 &self.mock_http
140 }
141
142 pub fn mock_http_mut(&mut self) -> &mut MockHttp {
144 &mut self.mock_http
145 }
146
147 pub fn dispatch_job(&mut self, job_type: &str, input: serde_json::Value) -> Uuid {
149 let job_id = Uuid::new_v4();
150 self.dispatched_jobs.push(DispatchedJob {
151 id: job_id,
152 job_type: job_type.to_string(),
153 input,
154 dispatched_at: Utc::now(),
155 status: JobStatus::Pending,
156 });
157 job_id
158 }
159
160 pub fn cancel_job(&mut self, job_id: Uuid) {
162 if let Some(job) = self.dispatched_jobs.iter_mut().find(|j| j.id == job_id) {
163 job.status = JobStatus::Cancelled;
164 }
165 }
166
167 pub fn dispatched_jobs(&self) -> &[DispatchedJob] {
169 &self.dispatched_jobs
170 }
171
172 pub fn job_dispatched(&self, job_type: &str) -> bool {
174 self.dispatched_jobs.iter().any(|j| j.job_type == job_type)
175 }
176
177 pub fn job_status(&self, job_id: Uuid) -> Option<JobStatus> {
179 self.dispatched_jobs
180 .iter()
181 .find(|j| j.id == job_id)
182 .map(|j| j.status)
183 }
184
185 pub fn complete_job(&mut self, job_id: Uuid) {
187 if let Some(job) = self.dispatched_jobs.iter_mut().find(|j| j.id == job_id) {
188 job.status = JobStatus::Completed;
189 }
190 }
191
192 pub fn run_jobs(&mut self) {
194 for job in &mut self.dispatched_jobs {
195 if job.status == JobStatus::Pending {
196 job.status = JobStatus::Completed;
197 }
198 }
199 }
200
201 pub fn start_workflow(&mut self, workflow_name: &str, input: serde_json::Value) -> Uuid {
203 let run_id = Uuid::new_v4();
204 self.started_workflows.push(StartedWorkflow {
205 run_id,
206 workflow_name: workflow_name.to_string(),
207 input,
208 started_at: Utc::now(),
209 status: WorkflowStatus::Created,
210 completed_steps: Vec::new(),
211 });
212 run_id
213 }
214
215 pub fn started_workflows(&self) -> &[StartedWorkflow] {
217 &self.started_workflows
218 }
219
220 pub fn workflow_status(&self, run_id: Uuid) -> Option<WorkflowStatus> {
222 self.started_workflows
223 .iter()
224 .find(|w| w.run_id == run_id)
225 .map(|w| w.status)
226 }
227
228 pub fn complete_workflow_step(&mut self, run_id: Uuid, step_name: &str) {
230 if let Some(workflow) = self
231 .started_workflows
232 .iter_mut()
233 .find(|w| w.run_id == run_id)
234 {
235 workflow.completed_steps.push(step_name.to_string());
236 }
237 }
238
239 pub fn complete_workflow(&mut self, run_id: Uuid) {
241 if let Some(workflow) = self
242 .started_workflows
243 .iter_mut()
244 .find(|w| w.run_id == run_id)
245 {
246 workflow.status = WorkflowStatus::Completed;
247 }
248 }
249
250 pub fn workflow_step_completed(&self, run_id: Uuid, step_name: &str) -> bool {
252 self.started_workflows
253 .iter()
254 .find(|w| w.run_id == run_id)
255 .map(|w| w.completed_steps.contains(&step_name.to_string()))
256 .unwrap_or(false)
257 }
258}
259
260pub struct TestContextBuilder {
262 config: TestConfig,
263 user_id: Option<Uuid>,
264 roles: Vec<String>,
265 custom_claims: HashMap<String, serde_json::Value>,
266 mock_http: MockHttp,
267}
268
269impl TestContextBuilder {
270 pub fn new() -> Self {
272 Self {
273 config: TestConfig::default(),
274 user_id: None,
275 roles: Vec::new(),
276 custom_claims: HashMap::new(),
277 mock_http: MockHttp::new(),
278 }
279 }
280
281 pub fn database_url(mut self, url: impl Into<String>) -> Self {
283 self.config.database_url = Some(url.into());
284 self
285 }
286
287 pub fn as_user(mut self, user_id: Uuid) -> Self {
289 self.user_id = Some(user_id);
290 self
291 }
292
293 pub fn with_roles(mut self, roles: Vec<String>) -> Self {
295 self.roles = roles;
296 self
297 }
298
299 pub fn with_claims(mut self, claims: HashMap<String, serde_json::Value>) -> Self {
301 self.custom_claims = claims;
302 self
303 }
304
305 pub fn with_logging(mut self, enabled: bool) -> Self {
307 self.config.logging = enabled;
308 self
309 }
310
311 pub fn mock_http(
313 mut self,
314 pattern: &str,
315 handler: impl Fn(&MockRequest) -> MockResponse + Send + Sync + 'static,
316 ) -> Self {
317 self.mock_http.add_mock(pattern, handler);
318 self
319 }
320
321 pub async fn build(self) -> Result<TestContext> {
323 let mut ctx = TestContext::with_config(self.config).await?;
324
325 if let Some(user_id) = self.user_id {
326 ctx.auth = AuthContext::authenticated(user_id, self.roles, self.custom_claims);
327 }
328
329 ctx.mock_http = self.mock_http;
330
331 Ok(ctx)
332 }
333}
334
335impl Default for TestContextBuilder {
336 fn default() -> Self {
337 Self::new()
338 }
339}
340
341#[cfg(test)]
342mod tests {
343 use super::*;
344
345 #[test]
346 fn test_context_builder() {
347 let builder = TestContextBuilder::new()
348 .as_user(Uuid::new_v4())
349 .with_logging(true);
350
351 assert!(builder.user_id.is_some());
352 assert!(builder.config.logging);
353 }
354
355 #[test]
356 fn test_context_without_db() {
357 let ctx = TestContext::new_without_db();
358 assert!(ctx.pool().is_none());
359 assert!(!ctx.auth().is_authenticated());
360 }
361
362 #[test]
363 fn test_job_dispatch() {
364 let mut ctx = TestContext::new_without_db();
365 let job_id = ctx.dispatch_job("send_email", serde_json::json!({"to": "test@example.com"}));
366
367 assert!(ctx.job_dispatched("send_email"));
368 assert_eq!(ctx.job_status(job_id), Some(JobStatus::Pending));
369
370 ctx.complete_job(job_id);
371 assert_eq!(ctx.job_status(job_id), Some(JobStatus::Completed));
372 }
373
374 #[test]
375 fn test_workflow_tracking() {
376 let mut ctx = TestContext::new_without_db();
377 let run_id = ctx.start_workflow(
378 "onboarding",
379 serde_json::json!({"email": "test@example.com"}),
380 );
381
382 assert_eq!(ctx.workflow_status(run_id), Some(WorkflowStatus::Created));
383
384 ctx.complete_workflow_step(run_id, "create_user");
385 assert!(ctx.workflow_step_completed(run_id, "create_user"));
386 assert!(!ctx.workflow_step_completed(run_id, "send_email"));
387
388 ctx.complete_workflow(run_id);
389 assert_eq!(ctx.workflow_status(run_id), Some(WorkflowStatus::Completed));
390 }
391
392 #[test]
393 fn test_run_jobs() {
394 let mut ctx = TestContext::new_without_db();
395 let job1 = ctx.dispatch_job("job1", serde_json::json!({}));
396 let job2 = ctx.dispatch_job("job2", serde_json::json!({}));
397
398 ctx.run_jobs();
399
400 assert_eq!(ctx.job_status(job1), Some(JobStatus::Completed));
401 assert_eq!(ctx.job_status(job2), Some(JobStatus::Completed));
402 }
403}