Skip to main content

forge_runtime/testing/
context.rs

1//! Test context for integration tests.
2
3use 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
17/// Test context for integration tests.
18///
19/// Provides an isolated testing environment with transaction-based
20/// isolation, mock HTTP support, and test utilities.
21pub struct TestContext {
22    /// Database pool (if connected).
23    pool: Option<sqlx::PgPool>,
24    /// HTTP mock.
25    mock_http: MockHttp,
26    /// Auth context.
27    auth: AuthContext,
28    /// Dispatched jobs for verification.
29    dispatched_jobs: Vec<DispatchedJob>,
30    /// Started workflows for verification.
31    started_workflows: Vec<StartedWorkflow>,
32}
33
34/// Record of a dispatched job.
35#[derive(Debug, Clone)]
36pub struct DispatchedJob {
37    /// Job ID.
38    pub id: Uuid,
39    /// Job type name.
40    pub job_type: String,
41    /// Job input.
42    pub input: serde_json::Value,
43    /// Dispatch time.
44    pub dispatched_at: DateTime<Utc>,
45    /// Current status (for test verification).
46    pub status: JobStatus,
47}
48
49/// Record of a started workflow.
50#[derive(Debug, Clone)]
51pub struct StartedWorkflow {
52    /// Run ID.
53    pub run_id: Uuid,
54    /// Workflow name.
55    pub workflow_name: String,
56    /// Input.
57    pub input: serde_json::Value,
58    /// Started time.
59    pub started_at: DateTime<Utc>,
60    /// Current status.
61    pub status: WorkflowStatus,
62    /// Completed steps.
63    pub completed_steps: Vec<String>,
64}
65
66impl TestContext {
67    /// Create a new test context (without database).
68    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    /// Create a new test context with database connection.
79    pub async fn new() -> Result<Self> {
80        let config = TestConfig::default();
81        Self::with_config(config).await
82    }
83
84    /// Create a test context with custom configuration.
85    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    /// Create a builder for more complex setup.
109    pub fn builder() -> TestContextBuilder {
110        TestContextBuilder::new()
111    }
112
113    /// Get the database pool.
114    pub fn pool(&self) -> Option<&sqlx::PgPool> {
115        self.pool.as_ref()
116    }
117
118    /// Get the auth context.
119    pub fn auth(&self) -> &AuthContext {
120        &self.auth
121    }
122
123    /// Get the user ID if authenticated.
124    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    /// Set the authenticated user.
133    pub fn set_user(&mut self, user_id: Uuid) {
134        self.auth = AuthContext::authenticated(user_id, vec![], HashMap::new());
135    }
136
137    /// Get the mock HTTP.
138    pub fn mock_http(&self) -> &MockHttp {
139        &self.mock_http
140    }
141
142    /// Get mutable mock HTTP.
143    pub fn mock_http_mut(&mut self) -> &mut MockHttp {
144        &mut self.mock_http
145    }
146
147    /// Dispatch a job for testing.
148    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    /// Cancel a job (for testing).
161    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    /// Get dispatched jobs.
168    pub fn dispatched_jobs(&self) -> &[DispatchedJob] {
169        &self.dispatched_jobs
170    }
171
172    /// Check if a job was dispatched.
173    pub fn job_dispatched(&self, job_type: &str) -> bool {
174        self.dispatched_jobs.iter().any(|j| j.job_type == job_type)
175    }
176
177    /// Get job status.
178    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    /// Mark a job as completed (for testing).
186    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    /// Run all pending jobs synchronously.
193    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    /// Start a workflow for testing.
202    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    /// Get started workflows.
216    pub fn started_workflows(&self) -> &[StartedWorkflow] {
217        &self.started_workflows
218    }
219
220    /// Get workflow status.
221    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    /// Mark a workflow step as completed.
229    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    /// Complete a workflow.
240    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    /// Check if a workflow step was completed.
251    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
260/// Builder for TestContext.
261pub 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    /// Create a new builder.
271    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    /// Set the database URL.
282    pub fn database_url(mut self, url: impl Into<String>) -> Self {
283        self.config.database_url = Some(url.into());
284        self
285    }
286
287    /// Set the authenticated user.
288    pub fn as_user(mut self, user_id: Uuid) -> Self {
289        self.user_id = Some(user_id);
290        self
291    }
292
293    /// Add roles.
294    pub fn with_roles(mut self, roles: Vec<String>) -> Self {
295        self.roles = roles;
296        self
297    }
298
299    /// Add custom claims.
300    pub fn with_claims(mut self, claims: HashMap<String, serde_json::Value>) -> Self {
301        self.custom_claims = claims;
302        self
303    }
304
305    /// Enable logging.
306    pub fn with_logging(mut self, enabled: bool) -> Self {
307        self.config.logging = enabled;
308        self
309    }
310
311    /// Add HTTP mock.
312    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    /// Build the test context.
322    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}