Refactor and add tests

* Clean up code
* Add tests