"""Tests for atm.state_machine.""" from __future__ import annotations import pytest from atm.state_machine import DotColor, State, StateMachine, Transition # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _drive(sm: StateMachine, events: list[tuple[DotColor, float]]) -> list[Transition]: return [sm.feed(color, ts) for color, ts in events] def _buy_sequence_to_primed(sm: StateMachine, start_ts: float = 1.0) -> None: """Drive sm from IDLE → PRIMED_BUY (does not fire).""" sm.feed("turquoise", start_ts) # → ARMED_BUY sm.feed("dark_green", start_ts + 1) # → PRIMED_BUY def _sell_sequence_to_primed(sm: StateMachine, start_ts: float = 1.0) -> None: """Drive sm from IDLE → PRIMED_SELL (does not fire).""" sm.feed("yellow", start_ts) # → ARMED_SELL sm.feed("dark_red", start_ts + 1) # → PRIMED_SELL # --------------------------------------------------------------------------- # 1. clean_buy # --------------------------------------------------------------------------- def test_clean_buy() -> None: sm = StateMachine() t1 = sm.feed("turquoise", 1.0) assert t1.prev == State.IDLE assert t1.next == State.ARMED_BUY assert t1.reason == "arm" assert sm.state == State.ARMED_BUY t2 = sm.feed("gray", 2.0) assert t2.prev == State.ARMED_BUY assert t2.next == State.ARMED_BUY assert t2.reason == "persist" assert sm.state == State.ARMED_BUY t3 = sm.feed("dark_green", 3.0) assert t3.prev == State.ARMED_BUY assert t3.next == State.PRIMED_BUY assert t3.reason == "prime" assert sm.state == State.PRIMED_BUY t4 = sm.feed("light_green", 4.0) assert t4.prev == State.PRIMED_BUY assert t4.next == State.IDLE assert t4.reason == "fire" assert t4.trigger == "BUY" assert t4.locked is False assert sm.state == State.IDLE # --------------------------------------------------------------------------- # 2. clean_sell (mirror of clean_buy) # --------------------------------------------------------------------------- def test_clean_sell() -> None: sm = StateMachine() t1 = sm.feed("yellow", 1.0) assert t1.prev == State.IDLE assert t1.next == State.ARMED_SELL assert t1.reason == "arm" t2 = sm.feed("gray", 2.0) assert t2.prev == State.ARMED_SELL assert t2.next == State.ARMED_SELL assert t2.reason == "persist" t3 = sm.feed("dark_red", 3.0) assert t3.prev == State.ARMED_SELL assert t3.next == State.PRIMED_SELL assert t3.reason == "prime" t4 = sm.feed("light_red", 4.0) assert t4.prev == State.PRIMED_SELL assert t4.next == State.IDLE assert t4.reason == "fire" assert t4.trigger == "SELL" assert t4.locked is False # --------------------------------------------------------------------------- # 3. cooled # --------------------------------------------------------------------------- def test_cooled() -> None: sm = StateMachine() _buy_sequence_to_primed(sm, start_ts=1.0) assert sm.state == State.PRIMED_BUY t = sm.feed("gray", 10.0) assert t.prev == State.PRIMED_BUY assert t.next == State.IDLE assert t.reason == "cooled" assert t.trigger is None # --------------------------------------------------------------------------- # 4. opposite_rearm from ARMED_BUY # --------------------------------------------------------------------------- def test_opposite_rearm_from_armed_buy() -> None: sm = StateMachine() sm.feed("turquoise", 1.0) # → ARMED_BUY assert sm.state == State.ARMED_BUY t = sm.feed("yellow", 2.0) assert t.prev == State.ARMED_BUY assert t.next == State.ARMED_SELL assert t.reason == "opposite_rearm" assert sm.state == State.ARMED_SELL # --------------------------------------------------------------------------- # 5. opposite_rearm from PRIMED_BUY # --------------------------------------------------------------------------- def test_opposite_rearm_from_primed_buy() -> None: sm = StateMachine() _buy_sequence_to_primed(sm, start_ts=1.0) assert sm.state == State.PRIMED_BUY t = sm.feed("yellow", 5.0) assert t.prev == State.PRIMED_BUY assert t.next == State.ARMED_SELL assert t.reason == "opposite_rearm" assert sm.state == State.ARMED_SELL # --------------------------------------------------------------------------- # 6. lockout_same_direction # --------------------------------------------------------------------------- def test_lockout_same_direction() -> None: sm = StateMachine(lockout_s=240) # First fire at t=100 _buy_sequence_to_primed(sm, start_ts=90.0) t_fire1 = sm.feed("light_green", 100.0) assert t_fire1.trigger == "BUY" assert t_fire1.locked is False # Re-prime _buy_sequence_to_primed(sm, start_ts=110.0) # Second fire at t=200 — inside lockout window (200 - 100 = 100 < 240) t_fire2 = sm.feed("light_green", 200.0) assert t_fire2.trigger == "BUY" assert t_fire2.locked is True assert t_fire2.next == State.IDLE # Re-prime again _buy_sequence_to_primed(sm, start_ts=300.0) # Third fire at t=341 — outside lockout (341 - 200 = 141 < 240, still locked) # Actually 341 - 100 would be outside but last fire is at t=200: # 341 - 200 = 141 < 240 → still locked # We need t > 200 + 240 = 440 _buy_sequence_to_primed(sm, start_ts=430.0) t_fire3 = sm.feed("light_green", 441.0) assert t_fire3.trigger == "BUY" assert t_fire3.locked is False def test_lockout_same_direction_boundary() -> None: """Spec requirement: fire BUY @ t=100; fire again @ t=341 → locked=False (241 >= 240).""" sm = StateMachine(lockout_s=240) _buy_sequence_to_primed(sm, start_ts=90.0) t_fire1 = sm.feed("light_green", 100.0) assert t_fire1.locked is False # Re-prime and fire just inside window: 339-100=239 < 240 → locked _buy_sequence_to_primed(sm, start_ts=110.0) t_locked = sm.feed("light_green", 339.0) assert t_locked.locked is True # last_fire is now 339. Re-prime and fire just outside: 580-339=241 >= 240 → unlocked _buy_sequence_to_primed(sm, start_ts=340.0) t_free = sm.feed("light_green", 580.0) assert t_free.locked is False # --------------------------------------------------------------------------- # 7. lockout_does_not_block_opposite # --------------------------------------------------------------------------- def test_lockout_does_not_block_opposite() -> None: sm = StateMachine(lockout_s=240) # Fire BUY at t=100 _buy_sequence_to_primed(sm, start_ts=90.0) sm.feed("light_green", 100.0) # Drive SELL sequence — opposite direction must not be locked _sell_sequence_to_primed(sm, start_ts=110.0) t_sell = sm.feed("light_red", 200.0) assert t_sell.trigger == "SELL" assert t_sell.locked is False # --------------------------------------------------------------------------- # 8. phase_skip from ARMED_BUY (light_green without priming) # --------------------------------------------------------------------------- def test_phase_skip_armed_buy() -> None: sm = StateMachine() sm.feed("turquoise", 1.0) # → ARMED_BUY assert sm.state == State.ARMED_BUY t = sm.feed("light_green", 2.0) assert t.prev == State.ARMED_BUY assert t.next == State.IDLE assert t.reason == "phase_skip" assert t.trigger is None # --------------------------------------------------------------------------- # 9. noise_from_idle # --------------------------------------------------------------------------- def test_noise_from_idle() -> None: sm = StateMachine() t = sm.feed("dark_green", 1.0) assert t.prev == State.IDLE assert t.next == State.IDLE assert t.reason == "noise" assert t.trigger is None assert sm.state == State.IDLE # --------------------------------------------------------------------------- # 10. refresh_arm_ts # --------------------------------------------------------------------------- def test_refresh_arm_ts() -> None: sm = StateMachine() sm.feed("turquoise", 1.0) # arm at t=1 t1 = sm.feed("turquoise", 5.0) # refresh at t=5 assert t1.prev == State.ARMED_BUY assert t1.next == State.ARMED_BUY assert t1.reason == "refresh" assert t1.arm_ts == 5.0 t2 = sm.feed("turquoise", 9.0) # refresh again at t=9 assert t2.arm_ts == 9.0 # --------------------------------------------------------------------------- # 11. exhaustive — parameterize over every (state, color) pair # --------------------------------------------------------------------------- ALL_STATES = list(State) ALL_COLORS: list[DotColor] = [ "turquoise", "yellow", "dark_green", "dark_red", "light_green", "light_red", "gray", ] FIRE_DIRECTIONS: dict[str, str] = { State.PRIMED_BUY.value: "BUY", State.PRIMED_SELL.value: "SELL", } VALID_STATES = set(State) def _sm_in_state(target: State) -> StateMachine: """Return a fresh StateMachine already in the given state.""" sm = StateMachine() match target: case State.IDLE: pass case State.ARMED_BUY: sm.feed("turquoise", 1.0) case State.ARMED_SELL: sm.feed("yellow", 1.0) case State.PRIMED_BUY: sm.feed("turquoise", 1.0) sm.feed("dark_green", 2.0) case State.PRIMED_SELL: sm.feed("yellow", 1.0) sm.feed("dark_red", 2.0) assert sm.state == target, f"Setup failed: wanted {target}, got {sm.state}" return sm @pytest.mark.parametrize("state", ALL_STATES) @pytest.mark.parametrize("color", ALL_COLORS) def test_exhaustive(state: State, color: DotColor) -> None: sm = _sm_in_state(state) t = sm.feed(color, 10.0) # (a) resulting state is valid assert t.next in VALID_STATES, f"Invalid next state: {t.next}" # (b) reason is non-empty assert t.reason, f"Empty reason for ({state}, {color})" # (c) if fire, trigger matches direction if t.reason == "fire": expected_dir = FIRE_DIRECTIONS.get(state.value) assert t.trigger == expected_dir, ( f"Wrong trigger for fire from {state}: got {t.trigger}, expected {expected_dir}" )