Refactor to reduce duplication

Based on Clay's code review feedback