Refactor tests
Add FSM tests