Coverage for src/chuck_data/commands/wizard/state.py: 0%

81 statements  

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

1""" 

2Wizard state management for setup wizard. 

3""" 

4 

5from dataclasses import dataclass, field 

6from enum import Enum 

7from typing import Dict, List, Optional, Any 

8 

9 

10class WizardStep(Enum): 

11 """Steps in the setup wizard.""" 

12 

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" 

19 

20 

21class WizardAction(Enum): 

22 """Actions the wizard can take.""" 

23 

24 CONTINUE = "continue" 

25 RETRY = "retry" 

26 EXIT = "exit" 

27 COMPLETE = "complete" 

28 

29 

30@dataclass 

31class WizardState: 

32 """State of the setup wizard.""" 

33 

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 

41 

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 

57 

58 

59@dataclass 

60class StepResult: 

61 """Result of processing a wizard step.""" 

62 

63 success: bool 

64 message: str 

65 next_step: Optional[WizardStep] = None 

66 action: WizardAction = WizardAction.CONTINUE 

67 data: Optional[Dict[str, Any]] = None 

68 

69 

70class WizardStateMachine: 

71 """State machine for managing wizard flow.""" 

72 

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 } 

96 

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, []) 

100 

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 

107 

108 if result.action == WizardAction.EXIT: 

109 # Exit the wizard 

110 return state 

111 

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 

118 

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) 

124 

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}" 

134 

135 return state 

136 

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