adding tests
Adding first tests, and refactoring a bit