Refactorings & Fix test cases