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

1"""Wait operation processor for handling WAIT operation updates.""" 

2 

3from __future__ import annotations 

4 

5from datetime import datetime, timedelta, timezone 

6from typing import Any 

7 

8from ....models import ( 

9 Operation, 

10 OperationAction, 

11 OperationStatus, 

12 OperationUpdate, 

13 WaitDetails, 

14) 

15from .base import ( 

16 OperationProcessor, 

17) 

18from ...exceptions import ( 

19 InvalidParameterValueException, 

20) 

21from ..time_scale import scale_delay 

22 

23VALID_ACTIONS_FOR_WAIT = frozenset( 

24 [ 

25 OperationAction.START, 

26 OperationAction.CANCEL, 

27 ] 

28) 

29 

30 

31class WaitProcessor(OperationProcessor): 

32 """Processes WAIT operation updates with timer scheduling.""" 

33 

34 valid_actions = VALID_ACTIONS_FOR_WAIT 

35 _ALLOWED_STATUS_TO_CANCEL = frozenset( 

36 [ 

37 OperationStatus.STARTED, 

38 ] 

39 ) 

40 

41 @classmethod 

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

43 """Validate WAIT operation update.""" 

44 super().validate(current_state, update) 

45 

46 match update.action: 

47 case OperationAction.START: 

48 if current_state is not None: 

49 msg_wait_exists: str = "Cannot start a WAIT that already exist." 

50 raise InvalidParameterValueException(msg_wait_exists) 

51 case OperationAction.CANCEL: 

52 if ( 

53 current_state is None 

54 or current_state.status not in cls._ALLOWED_STATUS_TO_CANCEL 

55 ): 

56 msg_wait_cancel: str = "Cannot cancel a WAIT that does not exist or has already completed." 

57 raise InvalidParameterValueException(msg_wait_cancel) 

58 

59 def process( 

60 self, 

61 update: OperationUpdate, 

62 current_op: Operation | None, 

63 notifier: Any, 

64 execution_arn: str, 

65 ) -> Operation: 

66 """Process WAIT operation update with scheduler integration for timers.""" 

67 match update.action: 

68 case OperationAction.START: 

69 wait_seconds = ( 

70 update.wait_options.wait_seconds if update.wait_options else 0 

71 ) 

72 scaled_wait_seconds = scale_delay(wait_seconds) 

73 

74 scheduled_end_timestamp = datetime.now(timezone.utc) + timedelta( 

75 seconds=scaled_wait_seconds 

76 ) 

77 

78 # Create WaitDetails with scheduled timestamp 

79 wait_details = WaitDetails( 

80 scheduled_end_timestamp=scheduled_end_timestamp 

81 ) 

82 

83 # Create new operation with wait details 

84 wait_operation = Operation( 

85 operation_id=update.operation_id, 

86 operation_type=update.operation_type, 

87 status=OperationStatus.STARTED, 

88 parent_id=update.parent_id, 

89 name=update.name, 

90 start_timestamp=datetime.now(timezone.utc), 

91 end_timestamp=None, 

92 sub_type=update.sub_type, 

93 execution_details=None, 

94 context_details=None, 

95 step_details=None, 

96 wait_details=wait_details, 

97 callback_details=None, 

98 chained_invoke_details=None, 

99 ) 

100 

101 # Schedule wait timer to complete after delay 

102 notifier.schedule_wait_timer( 

103 execution_arn, update.operation_id, scaled_wait_seconds 

104 ) 

105 return wait_operation 

106 case OperationAction.CANCEL: 

107 # TODO: need to cancel the WAIT in the executor 

108 # TODO: increase sequence id 

109 return self._translate_update_to_operation( 

110 update=update, 

111 current_operation=current_op, 

112 status=OperationStatus.CANCELLED, 

113 ) 

114 case _: 

115 msg: str = "Invalid action for WAIT operation." 

116 

117 raise InvalidParameterValueException(msg)