refactoring
modified test code to green