Minor changes to the unit testing class