Refactored to match spec