Small code review fixes, added test