Cleanup, refactor and added tests.