Coverage for async-durable-execution/src/async_durable_execution/runner/local/__init__.py: 97%

1from __future__ import annotations 

2 

3import asyncio 

4import inspect 

5import time 

6from typing import Any, Callable 

7 

8from ...client import DurableServiceClient 

9from ...execution import ( 

10 DurableExecutionInvocationInput, 

11 InitialExecutionState, 

12 _bind_service_client_to_handler, 

13) 

14from ...models import ( 

15 CheckpointOutput, 

16 CheckpointUpdatedExecutionState, 

17 DurableExecutionInvocationOutput, 

18 ErrorObject, 

19 Operation, 

20 OperationUpdate, 

21 StateOutput, 

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 

47 

48__all__ = [ 

49 "DurableFunctionLocalTestRunner", 

50 "Event", 

51 "Executor", 

52 "InMemoryServiceClient", 

53 "InProcessInvoker", 

54 "Scheduler", 

55 "create_local_runner", 

56 "create_test_lambda_context", 

57] 

58 

59 

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 ) 

74 

75 

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 ): 

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 ) 

102 

103 self._service_client.bind_executor(self._executor) 

104 

105 async def __aenter__(self) -> DurableFunctionLocalTestRunner: 

106 if scheduler := getattr(self, "_scheduler", None): 

107 scheduler.start() 

108 return self 

109 

110 async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: 

111 self.close() 

112 

113 def close(self): 

114 self._scheduler.stop() 

115 

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 ) 

123 

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) 

128 

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) 

133 

134 async def send_callback_heartbeat(self, callback_id: str) -> None: 

135 self._executor.send_callback_heartbeat(callback_id=callback_id) 

136 

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 ) 

143 

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 ) 

160 

161 output: StartDurableExecutionOutput = self._executor.start_execution( 

162 start_input 

163 ) 

164 

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 

169 

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) 

174 

175 if not completed: 

176 msg_timeout: str = "Execution did not complete within timeout" 

177 

178 raise TimeoutError(msg_timeout) 

179 

180 execution: Execution = self._executor.get_execution(execution_arn) 

181 return DurableFunctionTestResult.create(execution=execution) 

182 

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() 

187 

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 

201 

202 await asyncio.sleep(self.poll_interval) 

203 

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) 

208 

209 

210class InProcessInvoker(Invoker): 

211 def __init__(self, handler: Callable, service_client: InMemoryServiceClient): 

212 self.handler = _bind_service_client_to_handler(handler, service_client) 

213 self.service_client = service_client 

214 

215 def create_invocation_input( 

216 self, 

217 *, 

218 start_input: StartDurableExecutionInput, # noqa: ARG002 

219 durable_execution_arn: str, 

220 checkpoint_token: str, 

221 operations: list[Operation], 

222 ) -> DurableExecutionInvocationInput: 

223 return DurableExecutionInvocationInput( 

224 durable_execution_arn=durable_execution_arn, 

225 checkpoint_token=checkpoint_token, 

226 initial_execution_state=InitialExecutionState( 

227 operations=operations, 

228 next_marker="", 

229 ), 

230 ) 

231 

232 async def invoke( 

233 self, 

234 function_name: str, # noqa: ARG002 

235 input: DurableExecutionInvocationInput, 

236 endpoint_url: str | None = None, # noqa: ARG002 

237 ) -> InvokeResponse: 

238 context = create_test_lambda_context() 

239 payload = input.to_json_dict() 

240 async_handler = getattr(self.handler, "_async_handler", None) 

241 handler_result = ( 

242 async_handler(payload, context) 

243 if inspect.iscoroutinefunction(async_handler) 

244 else self.handler(payload, context) 

245 ) 

246 if inspect.isawaitable(handler_result): 

247 handler_result = await handler_result 

248 output = DurableExecutionInvocationOutput.from_dict(handler_result) 

249 return InvokeResponse( 

250 invocation_output=output, request_id=context.aws_request_id 

251 ) 

252 

253 def update_endpoint(self, endpoint_url: str, region_name: str) -> None: 

254 """No-op for in-process invoker.""" 

255 

256 

257def create_test_lambda_context() -> LambdaContext: 

258 # Create client context as a dictionary, not as objects 

259 # LambdaContext.__init__ expects dictionaries and will create the objects internally 

260 client_context_dict = { 

261 "custom": {"test_key": "test_value"}, 

262 "env": {"platform": "test", "make": "test", "model": "test"}, 

263 "client": { 

264 "installation_id": "test-installation-123", 

265 "app_title": "TestApp", 

266 "app_version_name": "1.0.0", 

267 "app_version_code": "100", 

268 "app_package_name": "com.test.app", 

269 }, 

270 } 

271 

272 cognito_identity_dict = { 

273 "cognitoIdentityId": "test-cognito-identity-123", 

274 "cognitoIdentityPoolId": "us-west-2:test-pool-456", 

275 } 

276 

277 return LambdaContext( 

278 aws_request_id="test-invoke-12345", 

279 client_context=client_context_dict, 

280 identity=cognito_identity_dict, 

281 invoked_function_arn="arn:aws:lambda:us-west-2:123456789012:function:test-function", 

282 tenant_id="test-tenant-789", 

283 ) 

284 

285 

286class InMemoryServiceClient(DurableServiceClient): 

287 """An in-memory service client, that can replace the boto lambda service client.""" 

288 

289 def __init__(self, scheduler: Scheduler): 

290 self._scheduler = scheduler 

291 self._executor = None 

292 self._transformer = OperationTransformer() 

293 

294 def bind_executor(self, executor) -> None: 

295 """Bind the local executor that handles checkpoint side effects.""" 

296 self._executor = executor 

297 

298 def mock_invoke_result(self, function_name: str, result: object) -> None: 

299 """Register a local mock result for a chained invoke function.""" 

300 self._transformer.mock_invoke_result(function_name=function_name, result=result) 

301 

302 def process_checkpoint( 

303 self, 

304 checkpoint_token: str, 

305 updates: list[OperationUpdate], 

306 client_token: str | None, # noqa: ARG002 

307 ) -> CheckpointOutput: 

308 """Process checkpoint updates and return result with updated execution state.""" 

309 if self._executor is None: 

310 unbound_msg = "Local executor is not bound to the service client." 

311 raise InvalidParameterValueException(unbound_msg) 

312 

313 token: CheckpointToken = CheckpointToken.from_str(checkpoint_token) 

314 execution: Execution = self._executor.get_execution(token.execution_arn) 

315 

316 if execution.is_complete or token.token_sequence != execution.token_sequence: 

317 msg: str = "Invalid checkpoint token" 

318 raise InvalidParameterValueException(msg) 

319 

320 CheckpointValidator.validate_input( 

321 updates, execution, processors=self._transformer.processors 

322 ) 

323 

324 updated_operations, all_updates = self._transformer.process_updates( 

325 updates=updates, 

326 current_operations=execution.operations, 

327 runner=self._executor, 

328 execution_arn=token.execution_arn, 

329 ) 

330 

331 new_checkpoint_token = execution.get_new_checkpoint_token() 

332 execution.operations = updated_operations 

333 execution.updates.extend(all_updates) 

334 self._executor.set_execution(execution) 

335 

336 return CheckpointOutput( 

337 checkpoint_token=new_checkpoint_token, 

338 new_execution_state=CheckpointUpdatedExecutionState( 

339 operations=execution.get_navigable_operations(), next_marker=None 

340 ), 

341 ) 

342 

343 async def checkpoint( 

344 self, 

345 durable_execution_arn: str, # noqa: ARG002 

346 checkpoint_token: str, 

347 updates: list[OperationUpdate], 

348 client_token: str | None, 

349 ) -> CheckpointOutput: 

350 # durable_execution_arn is not used in in-memory testing 

351 return self.process_checkpoint(checkpoint_token, updates, client_token) 

352 

353 async def get_execution_state( 

354 self, 

355 durable_execution_arn: str, # noqa: ARG002 

356 checkpoint_token: str, 

357 next_marker: str, 

358 max_items: int = 1000, 

359 ) -> StateOutput: 

360 # durable_execution_arn is not used in in-memory testing 

361 if self._executor is None: 

362 msg = "Local executor is not bound to the service client." 

363 raise InvalidParameterValueException(msg) 

364 

365 token: CheckpointToken = CheckpointToken.from_str(checkpoint_token) 

366 execution: Execution = self._executor.get_execution(token.execution_arn) 

367 

368 # TODO: paging when size or max 

369 return StateOutput( 

370 operations=execution.get_navigable_operations(), next_marker=None 

371 )