328 lines
10 KiB
Python
328 lines
10 KiB
Python
"""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}"
|
|
)
|