Minor change

Minor change to the original unit test codes