Fixing some of the unit tests