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

1"""Context operation processor for handling CONTEXT operation updates.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any 

6 

7from ....models import ( 

8 Operation, 

9 OperationAction, 

10 OperationStatus, 

11 OperationUpdate, 

12) 

13from .base import ( 

14 OperationProcessor, 

15) 

16from ...exceptions import ( 

17 InvalidParameterValueException, 

18) 

19 

20 

21VALID_ACTIONS_FOR_CONTEXT = frozenset( 

22 [ 

23 OperationAction.START, 

24 OperationAction.FAIL, 

25 OperationAction.SUCCEED, 

26 ] 

27) 

28 

29 

30class ContextProcessor(OperationProcessor): 

31 """Processes CONTEXT operation updates for execution context management.""" 

32 

33 valid_actions = VALID_ACTIONS_FOR_CONTEXT 

34 _ALLOWED_STATUS_TO_CLOSE = frozenset( 

35 [ 

36 OperationStatus.STARTED, 

37 ] 

38 ) 

39 

40 @classmethod 

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

42 """Validate CONTEXT operation update.""" 

43 super().validate(current_state, update) 

44 

45 match update.action: 

46 case OperationAction.START: 

47 if current_state is not None: 

48 msg_context_exists: str = ( 

49 "Cannot start a CONTEXT that already exist." 

50 ) 

51 raise InvalidParameterValueException(msg_context_exists) 

52 case OperationAction.FAIL | OperationAction.SUCCEED: 

53 if ( 

54 current_state is not None 

55 and current_state.status not in cls._ALLOWED_STATUS_TO_CLOSE 

56 ): 

57 msg_context_close: str = "Invalid current CONTEXT state to close." 

58 raise InvalidParameterValueException(msg_context_close) 

59 if update.action == OperationAction.FAIL and update.payload is not None: 

60 msg_context_fail_payload: str = ( 

61 "Cannot provide a Payload for FAIL action." 

62 ) 

63 raise InvalidParameterValueException(msg_context_fail_payload) 

64 if ( 

65 update.action == OperationAction.SUCCEED 

66 and update.error is not None 

67 ): 

68 msg_context_succeed_error: str = ( 

69 "Cannot provide an Error for SUCCEED action." 

70 ) 

71 raise InvalidParameterValueException(msg_context_succeed_error) 

72 

73 def process( 

74 self, 

75 update: OperationUpdate, 

76 current_op: Operation | None, 

77 notifier: Any, # noqa: ARG002 

78 execution_arn: str, # noqa: ARG002 

79 ) -> Operation: 

80 """Process CONTEXT operation update for context state transitions.""" 

81 match update.action: 

82 case OperationAction.START: 

83 return self._translate_update_to_operation( 

84 update=update, 

85 current_operation=current_op, 

86 status=OperationStatus.STARTED, 

87 ) 

88 case OperationAction.SUCCEED: 

89 return self._translate_update_to_operation( 

90 update=update, 

91 current_operation=current_op, 

92 status=OperationStatus.SUCCEEDED, 

93 ) 

94 case OperationAction.FAIL: 

95 return self._translate_update_to_operation( 

96 update=update, 

97 current_operation=current_op, 

98 status=OperationStatus.FAILED, 

99 ) 

100 case _: 

101 msg: str = "Invalid action for CONTEXT operation." 

102 raise InvalidParameterValueException(msg)