Refactor test case
Minor refactor