Coverage for async_durable_execution/_runner/local/__init__.py: 96%
140 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1from __future__ import annotations
3import asyncio
4import inspect
5import time
6from typing import Any, Callable
8from ..._core import (
9 CheckpointError,
10 CheckpointOutput,
11 CheckpointUpdatedExecutionState,
12 DurableExecutionInvocationInput,
13 DurableExecutionInvocationOutput,
14 DurableServiceClient,
15 ErrorObject,
16 GetExecutionStateError,
17 InitialExecutionState,
18 Operation,
19 OperationUpdate,
20 StateOutput,
21 _bind_service_client_to_handler,
22)
23from ..exceptions import (
24 DurableFunctionsTestError,
25 InvalidParameterValueException,
26 ResourceNotFoundException,
27)
28from .execution import Execution
29from .executor import Executor
30from ..model import (
31 DurableFunctionTestResult,
32 InvokeResponse,
33 _get_callback_id_from_events,
34)
35from .model import (
36 CheckpointToken,
37 Invoker,
38 LambdaContext,
39 StartDurableExecutionInput,
40 StartDurableExecutionOutput,
41)
42from .processor import (
43 CheckpointValidator,
44 OperationTransformer,
45)
46from .scheduler import Event, Scheduler
48__all__ = [
49 "DurableFunctionLocalTestRunner",
50 "Event",
51 "Executor",
52 "InMemoryServiceClient",
53 "InProcessInvoker",
54 "Scheduler",
55 "create_local_runner",
56 "create_test_lambda_context",
57]
60def create_local_runner(
61 *,
62 handler: Callable,
63 poll_interval: float = 1.0,
64 input: Any = None, # noqa: A002
65 timeout: int = 60,
66) -> DurableFunctionLocalTestRunner:
67 """Create a configured local durable function runner."""
68 return DurableFunctionLocalTestRunner(
69 handler=handler,
70 poll_interval=poll_interval,
71 input=input,
72 timeout=timeout,
73 )
76class DurableFunctionLocalTestRunner:
77 def __init__(
78 self,
79 handler: Callable,
80 poll_interval: float = 1.0,
81 input: Any = None, # noqa: A002
82 timeout: int = 900,
83 function_name: str = "test-function",
84 execution_name: str = "execution-name",
85 account_id: str = "123456789012",
86 ) -> None:
87 self._scheduler: Scheduler = Scheduler()
88 self.mode = "local"
89 self.poll_interval = poll_interval
90 self._default_input = input
91 self._default_timeout = timeout
92 self._function_name = function_name
93 self._execution_name = execution_name
94 self._account_id = account_id
95 self._service_client = InMemoryServiceClient(scheduler=self._scheduler)
96 self._invoker = InProcessInvoker(handler, self._service_client)
97 self._executor = Executor(
98 scheduler=self._scheduler,
99 invoker=self._invoker,
100 service_client=self._service_client,
101 )
103 self._service_client.bind_executor(self._executor)
105 async def __aenter__(self) -> DurableFunctionLocalTestRunner:
106 if scheduler := getattr(self, "_scheduler", None):
107 scheduler.start()
108 return self
110 async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
111 self.close()
113 def close(self) -> None:
114 self._scheduler.stop()
116 async def run(
117 self,
118 ) -> DurableFunctionTestResult:
119 execution_arn = await self.run_async()
120 return await self.wait_for_result(
121 execution_arn=execution_arn, timeout=self._default_timeout
122 )
124 async def send_callback_success(
125 self, callback_id: str, result: bytes | None = None
126 ) -> None:
127 self._executor.send_callback_success(callback_id=callback_id, result=result)
129 async def send_callback_failure(
130 self, callback_id: str, error: ErrorObject | None = None
131 ) -> None:
132 self._executor.send_callback_failure(callback_id=callback_id, error=error)
134 async def send_callback_heartbeat(self, callback_id: str) -> None:
135 self._executor.send_callback_heartbeat(callback_id=callback_id)
137 def mock_invoke_result(self, function_name: str, result: Any) -> None:
138 """Register a local mock result for a chained invoke function."""
139 self._service_client.mock_invoke_result(
140 function_name=function_name,
141 result=result,
142 )
144 async def run_async(
145 self,
146 ) -> str:
147 self._scheduler.start()
148 start_input = StartDurableExecutionInput(
149 account_id=self._account_id,
150 function_name=self._function_name,
151 function_qualifier="$LATEST",
152 execution_name=self._execution_name,
153 execution_timeout_seconds=self._default_timeout,
154 execution_retention_period_days=7,
155 invocation_id="inv-12345678-1234-1234-1234-123456789012",
156 trace_fields={"trace_id": "abc123", "span_id": "def456"},
157 tenant_id="tenant-001",
158 input=self._default_input,
159 )
161 output: StartDurableExecutionOutput = self._executor.start_execution(
162 start_input
163 )
165 if output.execution_arn is None:
166 msg_arn: str = "Execution ARN must exist to run test."
167 raise DurableFunctionsTestError(msg_arn)
168 return output.execution_arn
170 async def wait_for_result(
171 self, execution_arn: str, timeout: int = 60
172 ) -> DurableFunctionTestResult:
173 completed = await self._executor.wait_until_complete(execution_arn, timeout)
175 if not completed:
176 msg_timeout: str = "Execution did not complete within timeout"
178 raise TimeoutError(msg_timeout)
180 execution: Execution = self._executor.get_execution(execution_arn)
181 return DurableFunctionTestResult.create(execution=execution)
183 async def wait_for_callback(
184 self, execution_arn: str, name: str | None = None, timeout: int = 60
185 ) -> str:
186 start_time = time.time()
188 while time.time() - start_time < timeout:
189 try:
190 history_response = self._executor.get_execution_history(execution_arn)
191 callback_id = _get_callback_id_from_events(
192 events=history_response.events, name=name
193 )
194 if callback_id:
195 return callback_id
196 except ResourceNotFoundException:
197 pass
198 except Exception as e:
199 msg = f"Failed to fetch execution history: {e}"
200 raise DurableFunctionsTestError(msg) from e
202 await asyncio.sleep(self.poll_interval)
204 # Timeout reached
205 elapsed = time.time() - start_time
206 msg = f"Callback was not available within {timeout}s (elapsed: {elapsed:.1f}s)."
207 raise TimeoutError(msg)
210class InProcessInvoker(Invoker):
211 def __init__(
212 self, handler: Callable, service_client: InMemoryServiceClient
213 ) -> None:
214 self.handler = _bind_service_client_to_handler(handler, service_client)
215 self.service_client = service_client
217 def create_invocation_input(
218 self,
219 *,
220 start_input: StartDurableExecutionInput, # noqa: ARG002
221 durable_execution_arn: str,
222 checkpoint_token: str,
223 operations: list[Operation],
224 ) -> DurableExecutionInvocationInput:
225 return DurableExecutionInvocationInput(
226 durable_execution_arn=durable_execution_arn,
227 checkpoint_token=checkpoint_token,
228 initial_execution_state=InitialExecutionState(
229 operations=operations,
230 next_marker="",
231 ),
232 )
234 async def invoke(
235 self,
236 function_name: str, # noqa: ARG002
237 input: DurableExecutionInvocationInput,
238 endpoint_url: str | None = None, # noqa: ARG002
239 ) -> InvokeResponse:
240 context = create_test_lambda_context()
241 payload = input.to_dict()
242 async_handler = getattr(self.handler, "_async_handler", None)
243 handler_result = (
244 async_handler(payload, context)
245 if inspect.iscoroutinefunction(async_handler)
246 else self.handler(payload, context)
247 )
248 if inspect.isawaitable(handler_result):
249 handler_result = await handler_result
250 output = DurableExecutionInvocationOutput.from_dict(handler_result)
251 return InvokeResponse(
252 invocation_output=output, request_id=context.aws_request_id
253 )
256def create_test_lambda_context() -> LambdaContext:
257 # Create client context as a dictionary, not as objects
258 # LambdaContext.__init__ expects dictionaries and will create the objects internally
259 client_context_dict = {
260 "custom": {"test_key": "test_value"},
261 "env": {"platform": "test", "make": "test", "model": "test"},
262 "client": {
263 "installation_id": "test-installation-123",
264 "app_title": "TestApp",
265 "app_version_name": "1.0.0",
266 "app_version_code": "100",
267 "app_package_name": "com.test.app",
268 },
269 }
271 cognito_identity_dict = {
272 "cognitoIdentityId": "test-cognito-identity-123",
273 "cognitoIdentityPoolId": "us-west-2:test-pool-456",
274 }
276 return LambdaContext(
277 aws_request_id="test-invoke-12345",
278 client_context=client_context_dict,
279 identity=cognito_identity_dict,
280 invoked_function_arn="arn:aws:lambda:us-west-2:123456789012:function:test-function",
281 tenant_id="test-tenant-789",
282 )
285class InMemoryServiceClient(DurableServiceClient):
286 """An in-memory service client, that can replace the boto lambda service client."""
288 def __init__(self, scheduler: Scheduler) -> None:
289 self._scheduler = scheduler
290 self._executor = None
291 self._transformer = OperationTransformer()
293 def bind_executor(self, executor) -> None:
294 """Bind the local executor that handles checkpoint side effects."""
295 self._executor = executor
297 def mock_invoke_result(self, function_name: str, result: object) -> None:
298 """Register a local mock result for a chained invoke function."""
299 self._transformer.mock_invoke_result(function_name=function_name, result=result)
301 def process_checkpoint(
302 self,
303 checkpoint_token: str,
304 updates: list[OperationUpdate],
305 client_token: str | None, # noqa: ARG002
306 ) -> CheckpointOutput:
307 """Process checkpoint updates and return result with updated execution state."""
308 if self._executor is None: 308 ↛ 309line 308 didn't jump to line 309 because the condition on line 308 was never true
309 unbound_msg = "Local executor is not bound to the service client."
310 raise InvalidParameterValueException(unbound_msg)
312 token: CheckpointToken = CheckpointToken.from_str(checkpoint_token)
313 execution: Execution = self._executor.get_execution(token.execution_arn)
315 if execution.is_complete or token.token_sequence != execution.token_sequence:
316 msg: str = "Invalid checkpoint token"
317 raise InvalidParameterValueException(msg)
319 CheckpointValidator.validate_input(
320 updates, execution, processors=self._transformer.processors
321 )
323 updated_operations, all_updates = self._transformer.process_updates(
324 updates=updates,
325 current_operations=execution.operations,
326 runner=self._executor,
327 execution_arn=token.execution_arn,
328 )
330 new_checkpoint_token = execution.get_new_checkpoint_token()
331 execution.operations = updated_operations
332 execution.updates.extend(all_updates)
333 self._executor.set_execution(execution)
335 return CheckpointOutput(
336 checkpoint_token=new_checkpoint_token,
337 new_execution_state=CheckpointUpdatedExecutionState(
338 operations=execution.get_navigable_operations(), next_marker=None
339 ),
340 )
342 async def checkpoint(
343 self,
344 durable_execution_arn: str, # noqa: ARG002
345 checkpoint_token: str | None,
346 updates: list[OperationUpdate],
347 client_token: str | None,
348 ) -> CheckpointOutput:
349 # durable_execution_arn is not used in in-memory testing
350 if not checkpoint_token:
351 msg = "Cannot checkpoint without a checkpoint token."
352 raise CheckpointError(msg)
353 return self.process_checkpoint(checkpoint_token, updates, client_token)
355 async def get_execution_state(
356 self,
357 durable_execution_arn: str, # noqa: ARG002
358 checkpoint_token: str | None,
359 next_marker: str,
360 max_items: int = 1000,
361 ) -> StateOutput:
362 # durable_execution_arn is not used in in-memory testing
363 if not checkpoint_token:
364 msg = "Cannot get execution state without a checkpoint token."
365 raise GetExecutionStateError(msg)
367 if self._executor is None: 367 ↛ 368line 367 didn't jump to line 368 because the condition on line 367 was never true
368 msg = "Local executor is not bound to the service client."
369 raise InvalidParameterValueException(msg)
371 token: CheckpointToken = CheckpointToken.from_str(checkpoint_token)
372 execution: Execution = self._executor.get_execution(token.execution_arn)
374 # TODO: paging when size or max
375 return StateOutput(
376 operations=execution.get_navigable_operations(), next_marker=None
377 )