Various fixes and unit tests