Corrected and updated unit tests