Unit-test refactor