Correcting unittest mistake