Tweak the unit tests