More test case coverage
Refactoring