Started to update the unit tests