Refactor and start adding unit tests