started to write test and small refactoring