Refactor test script