Schedule : make `STATES` public and use an `assert` in the `wait` method

We missed `wait` in the transition from argument errors -> asserts.