Coverage for src/commands/wizard/state.py: 78%
81 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"""
2Wizard state management for setup wizard.
3"""
5from dataclasses import dataclass, field
6from enum import Enum
7from typing import Dict, List, Optional, Any
10class WizardStep(Enum):
11 """Steps in the setup wizard."""
13 AMPERITY_AUTH = "amperity_auth"
14 WORKSPACE_URL = "workspace_url"
15 TOKEN_INPUT = "token_input"
16 MODEL_SELECTION = "model_selection"
17 USAGE_CONSENT = "usage_consent"
18 COMPLETE = "complete"
21class WizardAction(Enum):
22 """Actions the wizard can take."""
24 CONTINUE = "continue"
25 RETRY = "retry"
26 EXIT = "exit"
27 COMPLETE = "complete"
30@dataclass
31class WizardState:
32 """State of the setup wizard."""
34 current_step: WizardStep = WizardStep.AMPERITY_AUTH
35 workspace_url: Optional[str] = None
36 token: Optional[str] = None
37 models: List[Dict[str, Any]] = field(default_factory=list)
38 selected_model: Optional[str] = None
39 usage_consent: Optional[bool] = None
40 error_message: Optional[str] = None
42 def is_valid_for_step(self, step: WizardStep) -> bool:
43 """Check if current state is valid for the given step."""
44 if step == WizardStep.AMPERITY_AUTH:
45 return True
46 elif step == WizardStep.WORKSPACE_URL:
47 return True # Can always enter workspace URL step
48 elif step == WizardStep.TOKEN_INPUT:
49 return self.workspace_url is not None
50 elif step == WizardStep.MODEL_SELECTION:
51 return self.workspace_url is not None and self.token is not None
52 elif step == WizardStep.USAGE_CONSENT:
53 return True # Can skip to usage consent if no models available
54 elif step == WizardStep.COMPLETE:
55 return self.usage_consent is not None
56 return False
59@dataclass
60class StepResult:
61 """Result of processing a wizard step."""
63 success: bool
64 message: str
65 next_step: Optional[WizardStep] = None
66 action: WizardAction = WizardAction.CONTINUE
67 data: Optional[Dict[str, Any]] = None
70class WizardStateMachine:
71 """State machine for managing wizard flow."""
73 def __init__(self):
74 self.valid_transitions = {
75 WizardStep.AMPERITY_AUTH: [
76 WizardStep.WORKSPACE_URL,
77 WizardStep.AMPERITY_AUTH,
78 ],
79 WizardStep.WORKSPACE_URL: [
80 WizardStep.TOKEN_INPUT,
81 WizardStep.WORKSPACE_URL,
82 ],
83 WizardStep.TOKEN_INPUT: [
84 WizardStep.MODEL_SELECTION,
85 WizardStep.USAGE_CONSENT,
86 WizardStep.TOKEN_INPUT,
87 WizardStep.WORKSPACE_URL,
88 ],
89 WizardStep.MODEL_SELECTION: [
90 WizardStep.USAGE_CONSENT,
91 WizardStep.MODEL_SELECTION,
92 ],
93 WizardStep.USAGE_CONSENT: [WizardStep.COMPLETE, WizardStep.USAGE_CONSENT],
94 WizardStep.COMPLETE: [],
95 }
97 def can_transition(self, from_step: WizardStep, to_step: WizardStep) -> bool:
98 """Check if transition is valid."""
99 return to_step in self.valid_transitions.get(from_step, [])
101 def transition(self, state: WizardState, result: StepResult) -> WizardState:
102 """Apply step result to state and transition to next step."""
103 if not result.success and result.action == WizardAction.RETRY:
104 # Stay on current step for retry
105 state.error_message = result.message
106 return state
108 if result.action == WizardAction.EXIT:
109 # Exit the wizard
110 return state
112 # Set error message for failed steps, clear on successful steps
113 if result.success:
114 state.error_message = None
115 elif result.message:
116 # Preserve error message for failed steps that continue to next step
117 state.error_message = result.message
119 # Apply any data changes from the step result
120 if result.data:
121 for key, value in result.data.items():
122 if hasattr(state, key):
123 setattr(state, key, value)
125 # Transition to next step if specified and valid
126 if result.next_step and self.can_transition(
127 state.current_step, result.next_step
128 ):
129 if state.is_valid_for_step(result.next_step):
130 state.current_step = result.next_step
131 else:
132 # Invalid state for next step, set error
133 state.error_message = f"Invalid state for step {result.next_step.value}"
135 return state
137 def get_next_step(self, current_step: WizardStep, state: WizardState) -> WizardStep:
138 """Determine the natural next step based on current step and state."""
139 if current_step == WizardStep.AMPERITY_AUTH:
140 return WizardStep.WORKSPACE_URL
141 elif current_step == WizardStep.WORKSPACE_URL:
142 return WizardStep.TOKEN_INPUT
143 elif current_step == WizardStep.TOKEN_INPUT:
144 # Skip to usage consent if no models available
145 return (
146 WizardStep.MODEL_SELECTION if state.models else WizardStep.USAGE_CONSENT
147 )
148 elif current_step == WizardStep.MODEL_SELECTION:
149 return WizardStep.USAGE_CONSENT
150 elif current_step == WizardStep.USAGE_CONSENT:
151 return WizardStep.COMPLETE
152 else:
153 return WizardStep.COMPLETE