Refactor to remove more duplicate code