Small test refactor