Do minor test refactoring