Updating the unit tests