Finished up refactor and added unit tests