Refactor tests (#16)