refactor test slightly