Refactored the code

improved coverage