Small bug fixes and new unit tests