Test fix and code refactoring :)