added test. refactored code