Finishing TODOs and added unit tests