Refactor to make testing easier and update tests