More tests
Refactor, added tests