Code refactoring +improved test