Coverage for src/chuck_data/commands/setup_wizard.py: 0%

106 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 22:56 -0700

1""" 

2Refactored setup wizard command with clean architecture. 

3""" 

4 

5import logging 

6from typing import Optional, Any 

7 

8from ..clients.databricks import DatabricksAPIClient 

9from ..commands.base import CommandResult 

10from ..interactive_context import InteractiveContext 

11from ..command_registry import CommandDefinition 

12from ..ui.tui import get_console 

13 

14from ..commands.wizard import ( 

15 WizardStep, 

16 WizardState, 

17 WizardStateMachine, 

18 WizardAction, 

19 WizardRenderer, 

20 InputValidator, 

21 create_step, 

22) 

23 

24 

25class SetupWizardOrchestrator: 

26 """Orchestrates the setup wizard flow.""" 

27 

28 def __init__(self): 

29 self.state_machine = WizardStateMachine() 

30 self.validator = InputValidator() 

31 self.renderer = WizardRenderer(get_console()) 

32 self.context = InteractiveContext() 

33 

34 def start_wizard(self) -> CommandResult: 

35 """Start the setup wizard from the beginning.""" 

36 # Set up interactive context - use the TUI alias that the service recognizes 

37 self.context.set_active_context("/setup") 

38 

39 # Initialize state 

40 state = WizardState(current_step=WizardStep.AMPERITY_AUTH) 

41 

42 # Store state in context for persistence across calls 

43 self._save_state_to_context(state) 

44 

45 # Handle the first step (Amperity auth) 

46 return self._process_step(state, "") 

47 

48 def handle_interactive_input(self, input_text: str) -> CommandResult: 

49 """Handle interactive input for current step.""" 

50 # Load state from context 

51 state = self._load_state_from_context() 

52 

53 if not state: 

54 # Context lost, restart wizard 

55 return self.start_wizard() 

56 

57 return self._process_step(state, input_text, render_step=False) 

58 

59 def _process_step( 

60 self, state: WizardState, input_text: str, render_step: bool = True 

61 ) -> CommandResult: 

62 """Process a single wizard step.""" 

63 try: 

64 # Store the current step before processing (this will be the "previous" step after transition) 

65 previous_step = state.current_step 

66 

67 # Create step handler 

68 step = create_step(state.current_step, self.validator) 

69 

70 # Render the step UI only if requested (e.g., when starting, not when processing input) 

71 if render_step: 

72 step_number = self.renderer.get_step_number(state.current_step) 

73 self.renderer.render_step(step, state, step_number) 

74 

75 # Handle special case for Amperity auth (no input needed) 

76 if state.current_step == WizardStep.AMPERITY_AUTH and not input_text: 

77 result = step.handle_input("", state) 

78 else: 

79 # Handle user input 

80 result = step.handle_input(input_text, state) 

81 

82 # Apply result to state 

83 state = self.state_machine.transition(state, result) 

84 

85 # Handle different actions 

86 if result.action == WizardAction.EXIT: 

87 self._clear_context() 

88 return CommandResult( 

89 success=False, message="exit_interactive:" + result.message 

90 ) 

91 

92 elif result.action == WizardAction.COMPLETE: 

93 self._clear_context() 

94 self.renderer.render_completion() 

95 return CommandResult(success=True, message=result.message) 

96 

97 elif result.action == WizardAction.RETRY: 

98 # Stay in same step, save state and ask for input again 

99 self._save_state_to_context(state) 

100 

101 # Re-render the step to show the error message to the user 

102 step_number = self.renderer.get_step_number(state.current_step) 

103 self.renderer.render_step(step, state, step_number, clear_screen=False) 

104 

105 return CommandResult(success=False, message=result.message) 

106 

107 else: # CONTINUE 

108 # Save updated state and move to next step 

109 self._save_state_to_context(state) 

110 

111 # Check if we're done 

112 if state.current_step == WizardStep.COMPLETE: 

113 self._clear_context() 

114 self.renderer.render_completion() 

115 return CommandResult( 

116 success=True, message="Setup wizard completed successfully!" 

117 ) 

118 

119 # If we transitioned to a new step, render it immediately 

120 if result.next_step and result.next_step != WizardStep.COMPLETE: 

121 # Only clear screen if the step was successful AND we're moving forward 

122 should_clear = ( 

123 result.success 

124 and self._should_clear_screen_after_step(previous_step) 

125 and self._is_forward_progression( 

126 previous_step, result.next_step 

127 ) 

128 ) 

129 

130 # Update state to the new step before rendering 

131 state.current_step = result.next_step 

132 next_step = create_step(result.next_step, self.validator) 

133 next_step_number = self.renderer.get_step_number(result.next_step) 

134 self.renderer.render_step( 

135 next_step, state, next_step_number, clear_screen=should_clear 

136 ) 

137 # Indicate that we rendered content so TUI knows not to show prompt immediately 

138 return CommandResult( 

139 success=True, 

140 message="", # Empty message to avoid duplicate display 

141 data={"rendered_next_step": True}, 

142 ) 

143 

144 return CommandResult(success=True, message=result.message) 

145 

146 except Exception as e: 

147 logging.error(f"Error in setup wizard: {e}") 

148 self._clear_context() 

149 return CommandResult( 

150 success=False, error=e, message=f"Setup wizard error: {e}" 

151 ) 

152 

153 def _save_state_to_context(self, state: WizardState): 

154 """Save wizard state to interactive context.""" 

155 context_data = { 

156 "current_step": state.current_step.value, 

157 "workspace_url": state.workspace_url, 

158 "token": state.token, 

159 "models": state.models, 

160 "selected_model": state.selected_model, 

161 "usage_consent": state.usage_consent, 

162 "error_message": state.error_message, 

163 } 

164 

165 for key, value in context_data.items(): 

166 self.context.store_context_data("/setup", key, value) 

167 

168 def _load_state_from_context(self) -> Optional[WizardState]: 

169 """Load wizard state from interactive context.""" 

170 try: 

171 context_data = self.context.get_context_data("/setup") 

172 

173 if not context_data: 

174 return None 

175 

176 # Convert step value back to enum (handle both string and numeric formats) 

177 step_value = context_data.get( 

178 "current_step", WizardStep.AMPERITY_AUTH.value 

179 ) 

180 try: 

181 if isinstance(step_value, int): 

182 # Handle old numeric format from TUI 

183 step_map = { 

184 1: WizardStep.AMPERITY_AUTH, 

185 2: WizardStep.WORKSPACE_URL, # TUI used 2 for both URL and token 

186 3: WizardStep.MODEL_SELECTION, 

187 4: WizardStep.USAGE_CONSENT, 

188 } 

189 current_step = step_map.get(step_value, WizardStep.AMPERITY_AUTH) 

190 else: 

191 current_step = WizardStep(step_value) 

192 except (ValueError, TypeError): 

193 current_step = WizardStep.AMPERITY_AUTH 

194 

195 return WizardState( 

196 current_step=current_step, 

197 workspace_url=context_data.get("workspace_url"), 

198 token=context_data.get("token"), 

199 models=context_data.get("models", []), 

200 selected_model=context_data.get("selected_model"), 

201 usage_consent=context_data.get("usage_consent"), 

202 error_message=context_data.get("error_message"), 

203 ) 

204 

205 except Exception as e: 

206 logging.error(f"Error loading wizard state from context: {e}") 

207 return None 

208 

209 def _should_clear_screen_after_step(self, completed_step: WizardStep) -> bool: 

210 """Determine if screen should be cleared after successful completion of a step.""" 

211 # Clear screen after successful completion of steps 1, 3, and 4 

212 return completed_step in [ 

213 WizardStep.AMPERITY_AUTH, # Step 1 

214 WizardStep.TOKEN_INPUT, # Step 3 

215 WizardStep.MODEL_SELECTION, # Step 4 

216 ] 

217 

218 def _is_forward_progression( 

219 self, from_step: WizardStep, to_step: WizardStep 

220 ) -> bool: 

221 """Check if we're moving forward in the wizard (not going back due to errors).""" 

222 step_order = [ 

223 WizardStep.AMPERITY_AUTH, 

224 WizardStep.WORKSPACE_URL, 

225 WizardStep.TOKEN_INPUT, 

226 WizardStep.MODEL_SELECTION, 

227 WizardStep.USAGE_CONSENT, 

228 WizardStep.COMPLETE, 

229 ] 

230 

231 try: 

232 from_index = step_order.index(from_step) 

233 to_index = step_order.index(to_step) 

234 return to_index > from_index 

235 except ValueError: 

236 # If steps not found in order, assume it's not forward progression 

237 return False 

238 

239 def _clear_context(self): 

240 """Clear the wizard context.""" 

241 self.context.clear_active_context("/setup") 

242 

243 

244def handle_command( 

245 client: Optional[DatabricksAPIClient], interactive_input: str = None, **kwargs: Any 

246) -> CommandResult: 

247 """ 

248 Setup wizard command handler using the new architecture. 

249 

250 Args: 

251 client: API client instance (can be None) 

252 interactive_input: Optional user input when in interactive mode 

253 

254 Returns: 

255 CommandResult with setup status 

256 """ 

257 orchestrator = SetupWizardOrchestrator() 

258 

259 # Check if we're in interactive mode or starting fresh 

260 context = InteractiveContext() 

261 

262 if context.is_in_interactive_mode() and context.current_command == "/setup": 

263 # Handle interactive input 

264 if interactive_input is None: 

265 interactive_input = "" 

266 return orchestrator.handle_interactive_input(interactive_input) 

267 else: 

268 # Start new wizard 

269 return orchestrator.start_wizard() 

270 

271 

272# Command definition for registration in the command registry 

273 

274DEFINITION = CommandDefinition( 

275 name="setup-wizard", 

276 description="Interactive setup wizard for first-time configuration", 

277 handler=handle_command, 

278 parameters={}, 

279 required_params=[], 

280 tui_aliases=["/setup", "/wizard"], 

281 needs_api_client=True, 

282 visible_to_user=True, 

283 visible_to_agent=False, 

284 usage_hint="Example: /setup to start the interactive setup wizard", 

285 supports_interactive_input=True, 

286)