Fixed broken builds

Small code changes and added a unit test