Coverage for async-durable-execution/src/async_durable_execution/runner/local/processors/invoke.py: 100%

1"""Chained invoke operation processor for local runner mocks.""" 

2 

3from __future__ import annotations 

4 

5import json 

6from typing import Any 

7 

8from ....models import ( 

9 Operation, 

10 OperationAction, 

11 OperationStatus, 

12 OperationUpdate, 

13) 

14from .base import OperationProcessor 

15from ...exceptions import InvalidParameterValueException 

16 

17VALID_ACTIONS_FOR_INVOKE = frozenset( 

18 [ 

19 OperationAction.START, 

20 OperationAction.CANCEL, 

21 ] 

22) 

23 

24 

25class ChainedInvokeProcessor(OperationProcessor): 

26 """Processes CHAINED_INVOKE updates and applies configured local mock results.""" 

27 

28 valid_actions = VALID_ACTIONS_FOR_INVOKE 

29 _ALLOWED_STATUS_TO_CANCEL = frozenset( 

30 [ 

31 OperationStatus.STARTED, 

32 ] 

33 ) 

34 

35 def __init__(self) -> None: 

36 self._mock_results: dict[str, str | None] = {} 

37 

38 def mock_result(self, function_name: str, result: Any) -> None: 

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

40 self._mock_results[function_name] = json.dumps(result) 

41 

42 def process( 

43 self, 

44 update: OperationUpdate, 

45 current_op: Operation | None, 

46 notifier: Any, # noqa: ARG002 

47 execution_arn: str, # noqa: ARG002 

48 ) -> Operation: 

49 """Process CHAINED_INVOKE updates for local execution.""" 

50 if update.action is OperationAction.CANCEL: 

51 return self._translate_update_to_operation( 

52 update=update, 

53 current_operation=current_op, 

54 status=OperationStatus.CANCELLED, 

55 ) 

56 if update.action is not OperationAction.START: 

57 msg = "Invalid action for CHAINED_INVOKE operation." 

58 raise InvalidParameterValueException(msg) 

59 

60 status = OperationStatus.STARTED 

61 result = None 

62 if update.chained_invoke_options: 

63 function_name = update.chained_invoke_options.function_name 

64 if function_name in self._mock_results: 

65 status = OperationStatus.SUCCEEDED 

66 result = self._mock_results[function_name] 

67 

68 mocked_update = update 

69 if result is not None: 

70 mocked_update = OperationUpdate( 

71 operation_id=update.operation_id, 

72 parent_id=update.parent_id, 

73 operation_type=update.operation_type, 

74 sub_type=update.sub_type, 

75 action=update.action, 

76 name=update.name, 

77 payload=result, 

78 chained_invoke_options=update.chained_invoke_options, 

79 ) 

80 

81 return self._translate_update_to_operation( 

82 update=mocked_update, 

83 current_operation=current_op, 

84 status=status, 

85 ) 

86 

87 @classmethod 

88 def validate(cls, current_state: Operation | None, update: OperationUpdate) -> None: 

89 """Validate CHAINED_INVOKE operation update.""" 

90 super().validate(current_state, update) 

91 

92 match update.action: 

93 case OperationAction.START: 

94 if current_state is not None: 

95 msg_invoke_exists: str = ( 

96 "Cannot start an INVOKE that already exist." 

97 ) 

98 raise InvalidParameterValueException(msg_invoke_exists) 

99 case OperationAction.CANCEL: 

100 if ( 

101 current_state is None 

102 or current_state.status not in cls._ALLOWED_STATUS_TO_CANCEL 

103 ): 

104 msg_invoke_cancel: str = "Cannot cancel an INVOKE that does not exist or has already completed." 

105 raise InvalidParameterValueException(msg_invoke_cancel)