Refactor the code to add api-test generate support