Refactoring
Added tests