Refactor test classes