durable_lambda_testing/mock_context.rs
1//! MockDurableContext — pre-loaded step results for local testing.
2//!
3//! Implements FR37-FR38: create mock context with pre-loaded results,
4//! run tests without AWS credentials.
5
6use std::sync::Arc;
7
8use aws_sdk_lambda::types::{
9 CallbackDetails, ChainedInvokeDetails, Operation, OperationStatus, OperationType, StepDetails,
10};
11use durable_lambda_core::context::DurableContext;
12use durable_lambda_core::operation_id::OperationIdGenerator;
13
14use crate::mock_backend::{BatchCallCounter, CheckpointRecorder, MockBackend, OperationRecorder};
15
16/// Builder for creating a [`DurableContext`] with pre-loaded step results.
17///
18/// `MockDurableContext` generates a `DurableContext` in **Replaying** mode
19/// by pre-loading completed operations. When the handler calls `ctx.step()`,
20/// the pre-loaded results are returned without executing the closure.
21///
22/// Operation IDs are generated deterministically using the same blake2b
23/// algorithm as the core engine, ensuring the nth `with_step_result` call
24/// corresponds to the nth `ctx.step()` call.
25///
26/// # Examples
27///
28/// ```no_run
29/// # async fn example() {
30/// use durable_lambda_testing::prelude::*;
31///
32/// let (mut ctx, calls, _ops) = MockDurableContext::new()
33/// .with_step_result("validate", r#"{"valid": true}"#)
34/// .with_step_result("charge", r#"100"#)
35/// .build()
36/// .await;
37///
38/// // Steps replay cached results — closures are NOT executed
39/// let result: Result<serde_json::Value, String> = ctx.step("validate", || async {
40/// panic!("not executed during replay");
41/// }).await.unwrap();
42///
43/// assert_eq!(result.unwrap(), serde_json::json!({"valid": true}));
44///
45/// // Verify no checkpoints were made (pure replay)
46/// assert_no_checkpoints(&calls).await;
47/// # }
48/// ```
49pub struct MockDurableContext {
50 id_gen: OperationIdGenerator,
51 operations: Vec<Operation>,
52}
53
54impl MockDurableContext {
55 /// Create a new `MockDurableContext` builder.
56 ///
57 /// # Examples
58 ///
59 /// ```no_run
60 /// # async fn example() {
61 /// use durable_lambda_testing::prelude::*;
62 ///
63 /// let (mut ctx, calls, _ops) = MockDurableContext::new()
64 /// .with_step_result("my_step", r#""hello""#)
65 /// .build()
66 /// .await;
67 /// # }
68 /// ```
69 pub fn new() -> Self {
70 Self {
71 id_gen: OperationIdGenerator::new(None),
72 operations: Vec::new(),
73 }
74 }
75
76 /// Add a successful step result to the mock history.
77 ///
78 /// The `result_json` is a JSON string representing the step's return value.
79 /// It will be returned by `ctx.step()` during replay without executing
80 /// the closure.
81 ///
82 /// # Arguments
83 ///
84 /// * `_name` — Step name (for documentation; the operation ID is position-based)
85 /// * `result_json` — JSON string of the step result (e.g., `r#"{"valid": true}"#`)
86 ///
87 /// # Examples
88 ///
89 /// ```no_run
90 /// # async fn example() {
91 /// use durable_lambda_testing::prelude::*;
92 ///
93 /// let (mut ctx, _, _ops) = MockDurableContext::new()
94 /// .with_step_result("validate", r#"42"#)
95 /// .build()
96 /// .await;
97 ///
98 /// let result: Result<i32, String> = ctx.step("validate", || async {
99 /// panic!("not executed");
100 /// }).await.unwrap();
101 ///
102 /// assert_eq!(result.unwrap(), 42);
103 /// # }
104 /// ```
105 pub fn with_step_result(mut self, _name: &str, result_json: &str) -> Self {
106 let op_id = self.id_gen.next_id();
107 let op = Operation::builder()
108 .id(&op_id)
109 .r#type(OperationType::Step)
110 .status(OperationStatus::Succeeded)
111 .start_timestamp(aws_smithy_types::DateTime::from_secs(0))
112 .step_details(StepDetails::builder().result(result_json).build())
113 .build()
114 .unwrap_or_else(|e| panic!("failed to build mock Operation: {e}"));
115 self.operations.push(op);
116 self
117 }
118
119 /// Add a failed step result to the mock history.
120 ///
121 /// The step will replay as a typed error. The `error_type` is the type
122 /// name and `error_json` is the serialized error data.
123 ///
124 /// # Arguments
125 ///
126 /// * `_name` — Step name (for documentation; the operation ID is position-based)
127 /// * `error_type` — The error type name (e.g., `"my_crate::MyError"`)
128 /// * `error_json` — JSON string of the error data
129 ///
130 /// # Examples
131 ///
132 /// ```no_run
133 /// # async fn example() {
134 /// use durable_lambda_testing::prelude::*;
135 ///
136 /// let (mut ctx, _, _ops) = MockDurableContext::new()
137 /// .with_step_error("charge", "PaymentError", r#""insufficient_funds""#)
138 /// .build()
139 /// .await;
140 ///
141 /// let result: Result<i32, String> = ctx.step("charge", || async {
142 /// panic!("not executed");
143 /// }).await.unwrap();
144 ///
145 /// assert_eq!(result.unwrap_err(), "insufficient_funds");
146 /// # }
147 /// ```
148 pub fn with_step_error(mut self, _name: &str, error_type: &str, error_json: &str) -> Self {
149 let op_id = self.id_gen.next_id();
150 let error_obj = aws_sdk_lambda::types::ErrorObject::builder()
151 .error_type(error_type)
152 .error_data(error_json)
153 .build();
154 let op = Operation::builder()
155 .id(&op_id)
156 .r#type(OperationType::Step)
157 .status(OperationStatus::Failed)
158 .start_timestamp(aws_smithy_types::DateTime::from_secs(0))
159 .step_details(StepDetails::builder().error(error_obj).build())
160 .build()
161 .unwrap_or_else(|e| panic!("failed to build mock Operation: {e}"));
162 self.operations.push(op);
163 self
164 }
165
166 /// Add a completed wait to the mock history.
167 ///
168 /// Simulates a wait that has already completed (SUCCEEDED). During replay,
169 /// `ctx.wait()` will return `Ok(())` immediately.
170 ///
171 /// # Arguments
172 ///
173 /// * `_name` — Wait name (for documentation; the operation ID is position-based)
174 ///
175 /// # Examples
176 ///
177 /// ```no_run
178 /// # async fn example() {
179 /// use durable_lambda_testing::prelude::*;
180 ///
181 /// let (mut ctx, _, _ops) = MockDurableContext::new()
182 /// .with_step_result("validate", r#"42"#)
183 /// .with_wait("cooldown")
184 /// .with_step_result("charge", r#"100"#)
185 /// .build()
186 /// .await;
187 /// # }
188 /// ```
189 pub fn with_wait(mut self, _name: &str) -> Self {
190 let op_id = self.id_gen.next_id();
191 let op = Operation::builder()
192 .id(&op_id)
193 .r#type(OperationType::Wait)
194 .status(OperationStatus::Succeeded)
195 .start_timestamp(aws_smithy_types::DateTime::from_secs(0))
196 .build()
197 .unwrap_or_else(|e| panic!("failed to build mock Wait Operation: {e}"));
198 self.operations.push(op);
199 self
200 }
201
202 /// Add a completed callback to the mock history.
203 ///
204 /// Simulates a callback that has been signaled with success. During replay,
205 /// `ctx.create_callback()` will return a `CallbackHandle` with the given
206 /// `callback_id`, and `ctx.callback_result()` will return the deserialized
207 /// result.
208 ///
209 /// # Arguments
210 ///
211 /// * `_name` — Callback name (for documentation; the operation ID is position-based)
212 /// * `callback_id` — The server-generated callback ID
213 /// * `result_json` — JSON string of the callback result
214 ///
215 /// # Examples
216 ///
217 /// ```no_run
218 /// # async fn example() {
219 /// use durable_lambda_testing::prelude::*;
220 ///
221 /// let (mut ctx, _, _ops) = MockDurableContext::new()
222 /// .with_callback("approval", "cb-123", r#""approved""#)
223 /// .build()
224 /// .await;
225 /// # }
226 /// ```
227 pub fn with_callback(mut self, _name: &str, callback_id: &str, result_json: &str) -> Self {
228 let op_id = self.id_gen.next_id();
229 let cb_details = CallbackDetails::builder()
230 .callback_id(callback_id)
231 .result(result_json)
232 .build();
233 let op = Operation::builder()
234 .id(&op_id)
235 .r#type(OperationType::Callback)
236 .status(OperationStatus::Succeeded)
237 .start_timestamp(aws_smithy_types::DateTime::from_secs(0))
238 .callback_details(cb_details)
239 .build()
240 .unwrap_or_else(|e| panic!("failed to build mock Callback Operation: {e}"));
241 self.operations.push(op);
242 self
243 }
244
245 /// Add a completed invoke to the mock history.
246 ///
247 /// Simulates a chained invoke that has completed successfully. During replay,
248 /// `ctx.invoke()` will return the deserialized result.
249 ///
250 /// # Arguments
251 ///
252 /// * `_name` — Invoke name (for documentation; the operation ID is position-based)
253 /// * `result_json` — JSON string of the invoke result
254 ///
255 /// # Examples
256 ///
257 /// ```no_run
258 /// # async fn example() {
259 /// use durable_lambda_testing::prelude::*;
260 ///
261 /// let (mut ctx, _, _ops) = MockDurableContext::new()
262 /// .with_invoke("call_processor", r#"{"status":"ok"}"#)
263 /// .build()
264 /// .await;
265 /// # }
266 /// ```
267 pub fn with_invoke(mut self, _name: &str, result_json: &str) -> Self {
268 let op_id = self.id_gen.next_id();
269 let details = ChainedInvokeDetails::builder().result(result_json).build();
270 let op = Operation::builder()
271 .id(&op_id)
272 .r#type(OperationType::ChainedInvoke)
273 .status(OperationStatus::Succeeded)
274 .start_timestamp(aws_smithy_types::DateTime::from_secs(0))
275 .chained_invoke_details(details)
276 .build()
277 .unwrap_or_else(|e| panic!("failed to build mock ChainedInvoke Operation: {e}"));
278 self.operations.push(op);
279 self
280 }
281
282 /// Build the mock context, returning a `DurableContext` and checkpoint call recorder.
283 ///
284 /// The returned `DurableContext` starts in **Replaying** mode if any
285 /// operations were pre-loaded, or **Executing** mode if none were added.
286 ///
287 /// # Returns
288 ///
289 /// A tuple of:
290 /// - `DurableContext` — ready for use with `ctx.step(...)` etc.
291 /// - `Arc<Mutex<Vec<CheckpointCall>>>` — inspect checkpoint calls after test
292 ///
293 /// # Errors
294 ///
295 /// Returns [`DurableError`] if context construction fails (should not happen
296 /// with mock data).
297 ///
298 /// # Examples
299 ///
300 /// ```no_run
301 /// # async fn example() {
302 /// use durable_lambda_testing::prelude::*;
303 ///
304 /// let (mut ctx, calls, _ops) = MockDurableContext::new()
305 /// .with_step_result("step1", r#"true"#)
306 /// .build()
307 /// .await;
308 /// # }
309 /// ```
310 pub async fn build(self) -> (DurableContext, CheckpointRecorder, OperationRecorder) {
311 let (backend, calls, operations) = MockBackend::new("mock-token");
312
313 let ctx = DurableContext::new(
314 Arc::new(backend),
315 "arn:aws:lambda:us-east-1:000000000000:durable-execution/mock".to_string(),
316 "mock-checkpoint-token".to_string(),
317 self.operations,
318 None,
319 )
320 .await
321 .expect("MockDurableContext::build should not fail");
322
323 (ctx, calls, operations)
324 }
325
326 /// Build mock context and also return the batch checkpoint call counter.
327 ///
328 /// Use this when testing batch mode — the `BatchCallCounter` lets you
329 /// assert how many times `batch_checkpoint()` was called.
330 ///
331 /// # Examples
332 ///
333 /// ```no_run
334 /// # async fn example() {
335 /// use durable_lambda_testing::prelude::*;
336 ///
337 /// let (mut ctx, calls, _ops, batch_counter) = MockDurableContext::new()
338 /// .build_with_batch_counter()
339 /// .await;
340 ///
341 /// ctx.enable_batch_mode();
342 /// let _: Result<i32, String> = ctx.step("s1", || async { Ok(1) }).await.unwrap();
343 /// ctx.flush_batch().await.unwrap();
344 ///
345 /// assert_eq!(*batch_counter.lock().await, 1);
346 /// # }
347 /// ```
348 pub async fn build_with_batch_counter(
349 self,
350 ) -> (
351 DurableContext,
352 CheckpointRecorder,
353 OperationRecorder,
354 BatchCallCounter,
355 ) {
356 let (backend, calls, operations) = MockBackend::new("mock-token");
357 let batch_counter = backend.batch_call_counter();
358
359 let ctx = DurableContext::new(
360 Arc::new(backend),
361 "arn:aws:lambda:us-east-1:000000000000:durable-execution/mock".to_string(),
362 "mock-checkpoint-token".to_string(),
363 self.operations,
364 None,
365 )
366 .await
367 .expect("MockDurableContext::build_with_batch_counter should not fail");
368
369 (ctx, calls, operations, batch_counter)
370 }
371}
372
373impl Default for MockDurableContext {
374 fn default() -> Self {
375 Self::new()
376 }
377}
378
379#[cfg(test)]
380mod tests {
381 use super::*;
382 use std::sync::atomic::{AtomicBool, Ordering};
383
384 #[tokio::test]
385 async fn test_mock_context_replays_step_result() {
386 let (mut ctx, calls, _ops) = MockDurableContext::new()
387 .with_step_result("validate", r#"42"#)
388 .build()
389 .await;
390
391 let executed = Arc::new(AtomicBool::new(false));
392 let executed_clone = executed.clone();
393
394 let result: Result<i32, String> = ctx
395 .step("validate", move || {
396 let executed = executed_clone.clone();
397 async move {
398 executed.store(true, Ordering::SeqCst);
399 Ok(999) // should NOT be returned
400 }
401 })
402 .await
403 .unwrap();
404
405 assert_eq!(result.unwrap(), 42);
406 assert!(
407 !executed.load(Ordering::SeqCst),
408 "closure should not execute during replay"
409 );
410
411 // No checkpoints should be made during pure replay
412 let captured = calls.lock().await;
413 assert_eq!(captured.len(), 0);
414 }
415
416 #[tokio::test]
417 async fn test_mock_context_replays_multiple_steps() {
418 let (mut ctx, calls, _ops) = MockDurableContext::new()
419 .with_step_result("step1", r#""hello""#)
420 .with_step_result("step2", r#""world""#)
421 .build()
422 .await;
423
424 let r1: Result<String, String> = ctx
425 .step("step1", || async { panic!("not executed") })
426 .await
427 .unwrap();
428 assert_eq!(r1.unwrap(), "hello");
429
430 let r2: Result<String, String> = ctx
431 .step("step2", || async { panic!("not executed") })
432 .await
433 .unwrap();
434 assert_eq!(r2.unwrap(), "world");
435
436 let captured = calls.lock().await;
437 assert_eq!(captured.len(), 0);
438 }
439
440 #[tokio::test]
441 async fn test_mock_context_replays_step_error() {
442 let (mut ctx, _calls, _ops) = MockDurableContext::new()
443 .with_step_error("charge", "PaymentError", r#""insufficient_funds""#)
444 .build()
445 .await;
446
447 let result: Result<i32, String> = ctx
448 .step("charge", || async { panic!("not executed") })
449 .await
450 .unwrap();
451
452 assert_eq!(result.unwrap_err(), "insufficient_funds");
453 }
454
455 #[tokio::test]
456 async fn test_mock_context_executing_mode_when_empty() {
457 let (ctx, _calls, _ops) = MockDurableContext::new().build().await;
458
459 assert!(!ctx.is_replaying());
460 assert_eq!(
461 ctx.execution_mode(),
462 durable_lambda_core::types::ExecutionMode::Executing
463 );
464 }
465
466 #[tokio::test]
467 async fn test_mock_context_replaying_mode_with_operations() {
468 let (ctx, _calls, _ops) = MockDurableContext::new()
469 .with_step_result("step1", r#"1"#)
470 .build()
471 .await;
472
473 assert!(ctx.is_replaying());
474 assert_eq!(
475 ctx.execution_mode(),
476 durable_lambda_core::types::ExecutionMode::Replaying
477 );
478 }
479
480 #[tokio::test]
481 async fn test_mock_context_no_aws_credentials_needed() {
482 // This test proves the mock works without any AWS env vars
483 // by simply running successfully
484 let (mut ctx, _calls, _ops) = MockDurableContext::new()
485 .with_step_result("test", r#"true"#)
486 .build()
487 .await;
488
489 let result: Result<bool, String> = ctx
490 .step("test", || async { panic!("not executed") })
491 .await
492 .unwrap();
493
494 assert!(result.unwrap());
495 }
496}