Refactor and fix current test cases.