Test class refactor