solve bug and remove the test code