test fix and refactoring