Improving test suite and some refactoring
refactoring