Refactor and start adding tests