Coverage for src/commands/setup_wizard.py: 70%
106 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
« 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"""
5import logging
6from typing import Optional, Any
8from src.clients.databricks import DatabricksAPIClient
9from src.commands.base import CommandResult
10from src.interactive_context import InteractiveContext
11from src.command_registry import CommandDefinition
12from src.ui.tui import get_console
14from src.commands.wizard import (
15 WizardStep,
16 WizardState,
17 WizardStateMachine,
18 WizardAction,
19 WizardRenderer,
20 InputValidator,
21 create_step,
22)
25class SetupWizardOrchestrator:
26 """Orchestrates the setup wizard flow."""
28 def __init__(self):
29 self.state_machine = WizardStateMachine()
30 self.validator = InputValidator()
31 self.renderer = WizardRenderer(get_console())
32 self.context = InteractiveContext()
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")
39 # Initialize state
40 state = WizardState(current_step=WizardStep.AMPERITY_AUTH)
42 # Store state in context for persistence across calls
43 self._save_state_to_context(state)
45 # Handle the first step (Amperity auth)
46 return self._process_step(state, "")
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()
53 if not state:
54 # Context lost, restart wizard
55 return self.start_wizard()
57 return self._process_step(state, input_text, render_step=False)
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
67 # Create step handler
68 step = create_step(state.current_step, self.validator)
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)
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)
82 # Apply result to state
83 state = self.state_machine.transition(state, result)
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 )
92 elif result.action == WizardAction.COMPLETE:
93 self._clear_context()
94 self.renderer.render_completion()
95 return CommandResult(success=True, message=result.message)
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)
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)
105 return CommandResult(success=False, message=result.message)
107 else: # CONTINUE
108 # Save updated state and move to next step
109 self._save_state_to_context(state)
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 )
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 )
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 )
144 return CommandResult(success=True, message=result.message)
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 )
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 }
165 for key, value in context_data.items():
166 self.context.store_context_data("/setup", key, value)
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")
173 if not context_data:
174 return None
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
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 )
205 except Exception as e:
206 logging.error(f"Error loading wizard state from context: {e}")
207 return None
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 ]
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 ]
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
239 def _clear_context(self):
240 """Clear the wizard context."""
241 self.context.clear_active_context("/setup")
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.
250 Args:
251 client: API client instance (can be None)
252 interactive_input: Optional user input when in interactive mode
254 Returns:
255 CommandResult with setup status
256 """
257 orchestrator = SetupWizardOrchestrator()
259 # Check if we're in interactive mode or starting fresh
260 context = InteractiveContext()
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()
272# Command definition for registration in the command registry
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)