More tests and refactoring
More tests and refactoring