minor change

minor change to fix the unit test code