Refactored and added tests
More tests