Refactoring
Fixed test