More functionality and unit tests