Refactored
Added test