Fixes and updates based on unit testing