hypothesis_fsm_example.py

import asyncio
import enum
import os
import unittest
from hypothesis import strategies as st
from hypothesis.stateful import Bundle, RuleBasedStateMachine, initialize, rule


class E(enum.Enum):
    SET = 1
    CLEAR = 2
    JOB = 3
    JOBDONE = 4


class S(enum.Enum):
    BLANK = 1
    ASSIGNED = 2
    BUSY = 3


machine = {
    S.BLANK: {E.SET: S.ASSIGNED},
    S.ASSIGNED: {E.CLEAR: S.BLANK, E.JOB: S.BUSY},
    S.BUSY: {E.JOBDONE: S.ASSIGNED},
}


class Director:
    state = S.BLANK

    def handle_event(self, event: E) -> S:
        match self.state, event:
            case S.BLANK, E.SET:
                print("got assignment")
                return S.ASSIGNED
            case S.BLANK, _:
                raise Exception
            case S.ASSIGNED, E.CLEAR:
                print("unassigned")
                return S.BLANK
            case S.ASSIGNED, E.JOB:
                print("got work")
                return S.BUSY
            case S.ASSIGNED, _:
                raise Exception
            case S.BUSY, E.JOBDONE:
                print("work done")
                return S.ASSIGNED
            case S.BUSY, E.CLEAR if os.getenv("BUG"):
                # example bug
                print("incorrect")
                return S.BLANK
            case _:
                raise Exception

    def step(self, event: E) -> None:
        self.state = self.handle_event(event)


class StateMachineTest(RuleBasedStateMachine):
    def __init__(self) -> None:
        super().__init__()
        self.director = Director()
        self.current_state = S.BLANK

    state = Bundle("state")

    @initialize(target=state)
    def initial_state(self) -> S:
        return S.BLANK

    @rule(target=state, event=st.sampled_from(list(E)))
    def transition(self, event: E) -> None:
        next_state = machine[self.current_state].get(event)
        try:
            self.director.step(event)
        except:
            assert next_state is None, "unexpected error"
        else:
            assert next_state, "director wrongly accepted invalid state"
            self.current_state = next_state
            print(self.director.state, self.current_state)
            assert self.director.state == next_state, "wrong state"


TestDirector = StateMachineTest.TestCase
if __name__ == "__main__":
    unittest.main()