Write tests

* Add tests to cover all current functionality
* Update TODOs