Refactor some test code