Little refactor on tests