Coverage for src/dataknobs_fsm/core/context_factory.py: 83%

1"""Factory for creating and configuring ExecutionContext instances.

2

3 This module provides a centralized factory for creating execution contexts,

4 eliminating duplication between Simple and Advanced APIs.

5"""

6

7 from typing import Any, Dict, Union

8 from dataknobs_data import Record

9

10 from ..core.fsm import FSM

11 from ..core.state import StateInstance, StateDefinition, StateType

12 from ..core.modes import ProcessingMode, TransactionMode

13 from ..execution.context import ExecutionContext

14 from ..resources.manager import ResourceManager

15

16

17 class ContextFactory:

18 """Factory for creating and configuring ExecutionContext instances."""

19

20 @staticmethod

21 def create_context(

22 fsm: FSM,

23 data: Union[Dict[str, Any], Record],

24 initial_state: str | None = None,

25 data_mode: ProcessingMode = ProcessingMode.SINGLE,

26 transaction_mode: TransactionMode = TransactionMode.NONE,

27 resources: Dict[str, Any] | None = None,

28 resource_manager: ResourceManager | None = None

29 ) -> ExecutionContext:

30 """Create and configure an ExecutionContext.

31

32 Args:

33 fsm: The FSM instance

34 data: Input data as dict or Record

35 initial_state: Optional initial state name

36 data_mode: Processing mode (single, batch, stream)

37 transaction_mode: Transaction mode

38 resources: Resource configuration dict

39 resource_manager: Optional resource manager instance

40

41 Returns:

42 Configured ExecutionContext ready for execution

43 """

44 # Create base context

45 context = ExecutionContext(

46 data_mode=data_mode,

47 transaction_mode=transaction_mode,

48 resources=resources or {}

49 )

50

51 # Convert and set data

52 if isinstance(data, Record):

53 context.data = data.to_dict()

54 else:

55 context.data = data

56

57 # Determine initial state

58 state_name, state_def = ContextFactory._resolve_initial_state(

59 fsm, initial_state

60 )

61

62 # Set state in context

63 context.current_state = state_name

64 if state_def:

65 context.current_state_instance = StateInstance(

66 definition=state_def,

67 data=context.data.copy() if isinstance(context.data, dict) else {}

68 )

69

70 # Add resource manager reference if provided

71 if resource_manager:

72 context.resource_manager = resource_manager

73

74 return context

75

76 @staticmethod

77 def _resolve_initial_state(

78 fsm: FSM,

79 state_name: str | None = None

80 ) -> tuple[str, StateDefinition | None]:

81 """Resolve the initial state for execution.

82

83 Args:

84 fsm: The FSM instance

85 state_name: Optional requested state name

86

87 Returns:

88 Tuple of (state_name, state_definition)

89 """

90 state_def = None

91

92 if state_name:

93 # Try to get the requested state

94 if hasattr(fsm, 'get_state'):

95 try:

96 state_def = fsm.get_state(state_name)

97 except (KeyError, AttributeError):

98 pass

99

100 # If not found, search in networks

101 if not state_def and hasattr(fsm, 'networks'):

102 for network in fsm.networks.values():

103 if hasattr(network, 'states') and state_name in network.states:

104 state_def = network.states[state_name]

105 break

106 else:

107 # Find the start state

108 if hasattr(fsm, 'get_start_state'):

109 state_def = fsm.get_start_state()

110 if state_def:

111 state_name = state_def.name

112

113 # Fallback: search networks for start state

114 if not state_def and hasattr(fsm, 'networks'):

115 for network in fsm.networks.values():

116 if hasattr(network, 'states'):

117 for state in network.states.values():

118 if hasattr(state, 'is_start_state') and state.is_start_state():

119 state_def = state

120 state_name = state.name

121 break

122 if state_def:

123 break

124

125 # Final fallback

126 if not state_name:

127 state_name = 'start'

128

129 # Create minimal state definition if needed

130 if not state_def and state_name:

131 state_def = StateDefinition(

132 name=state_name,

133 type=StateType.START if state_name == 'start' else StateType.NORMAL

134 )

135

136 return state_name, state_def

137

138 @staticmethod

139 def create_batch_context(

140 fsm: FSM,

141 batch_data: list,

142 data_mode: ProcessingMode = ProcessingMode.BATCH,

143 resources: Dict[str, Any] | None = None

144 ) -> ExecutionContext:

145 """Create a context for batch processing.

146

147 Args:

148 fsm: The FSM instance

149 batch_data: List of data items to process

150 data_mode: Processing mode (defaults to BATCH)

151 resources: Resource configuration

152

153 Returns:

154 ExecutionContext configured for batch processing

155 """

156 context = ExecutionContext(

157 data_mode=data_mode,

158 resources=resources or {}

159 )

160

161 # Set batch data

162 context.batch_data = batch_data

163

164 # Find start state

165 state_name, state_def = ContextFactory._resolve_initial_state(fsm, None)

166 context.current_state = state_name

167

168 if state_def:

169 context.current_state_instance = StateInstance(

170 definition=state_def,

171 data={}

172 )

173

174 return context

175

176 @staticmethod

177 def create_stream_context(

178 fsm: FSM,

179 stream_context,

180 resources: Dict[str, Any] | None = None

181 ) -> ExecutionContext:

182 """Create a context for stream processing.

183

184 Args:

185 fsm: The FSM instance

186 stream_context: Stream context object

187 resources: Resource configuration

188

189 Returns:

190 ExecutionContext configured for stream processing

191 """

192 context = ExecutionContext(

193 data_mode=ProcessingMode.STREAM,

194 resources=resources or {},

195 stream_context=stream_context

196 )

197

198 # Find start state

199 state_name, state_def = ContextFactory._resolve_initial_state(fsm, None)

200 context.current_state = state_name

201

202 if state_def:

203 context.current_state_instance = StateInstance(

204 definition=state_def,

205 data={}

206 )

207

208 return context