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

1"""Callback operation processor for handling CALLBACK operation updates.""" 

2 

3from __future__ import annotations 

4 

5import datetime 

6from typing import Any 

7 

8from ....models import ( 

9 CallbackDetails, 

10 CallbackOptions, 

11 Operation, 

12 OperationAction, 

13 OperationStatus, 

14 OperationType, 

15 OperationUpdate, 

16) 

17from .base import OperationProcessor 

18from ...exceptions import InvalidParameterValueException 

19from ..model import CallbackToken 

20 

21 

22VALID_ACTIONS_FOR_CALLBACK = frozenset( 

23 [ 

24 OperationAction.START, 

25 ] 

26) 

27 

28 

29class CallbackProcessor(OperationProcessor): 

30 """Processes CALLBACK operation updates with activity scheduling.""" 

31 

32 valid_actions = VALID_ACTIONS_FOR_CALLBACK 

33 

34 @classmethod 

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

36 """Validate CALLBACK operation update.""" 

37 super().validate(current_state, update) 

38 if current_state is not None: 

39 msg_callback_exists: str = "Cannot start a CALLBACK that already exist." 

40 raise InvalidParameterValueException(msg_callback_exists) 

41 

42 def process( 

43 self, 

44 update: OperationUpdate, 

45 current_op: Operation | None, 

46 notifier: Any, 

47 execution_arn: str, 

48 ) -> Operation: 

49 """Process CALLBACK operation update with scheduler integration for activities.""" 

50 match update.action: 

51 case OperationAction.START: 

52 callback_token: CallbackToken = CallbackToken( 

53 execution_arn=execution_arn, 

54 operation_id=update.operation_id, 

55 ) 

56 

57 callback_id: str = callback_token.to_str() 

58 

59 callback_details: CallbackDetails | None = ( 

60 CallbackDetails( 

61 callback_id=callback_id, 

62 result=update.payload, 

63 error=update.error, 

64 ) 

65 if update.operation_type == OperationType.CALLBACK 

66 else None 

67 ) 

68 

69 status: OperationStatus = OperationStatus.STARTED 

70 

71 start_time: datetime.datetime | None = self._get_start_time(current_op) 

72 

73 end_time: datetime.datetime | None = self._get_end_time( 

74 current_op, status 

75 ) 

76 

77 operation: Operation = Operation( 

78 operation_id=update.operation_id, 

79 parent_id=update.parent_id, 

80 name=update.name, 

81 start_timestamp=start_time, 

82 end_timestamp=end_time, 

83 operation_type=update.operation_type, 

84 status=status, 

85 sub_type=update.sub_type, 

86 callback_details=callback_details, 

87 ) 

88 callback_options: CallbackOptions | None = update.callback_options 

89 

90 notifier.schedule_callback_timeouts( 

91 execution_arn, callback_options, callback_id 

92 ) 

93 return operation 

94 case _: 

95 msg: str = "Invalid action for CALLBACK operation." 

96 raise InvalidParameterValueException(msg)