Refactoring code for third test
Third test is green