Refactor test classes to be more consistent