Refactor tests (again)