Minor refactor to reduce scope of bug