Refactor the test code.