Change unit test code