Fixing test
Refactoring