Refactoring tests a bit more, adding state tests