Move custom assertion method.

Defining the ```assert_correct_state``` within a shoulda context
meant that in some cases the method was not found. Moving it to a
wider context (the test class) fixes this.