Increased test coverage
Test refactor