Tests added and some refactoring