Refactored code to enable testing