Bug fix, small refactor.
Formatting
Adding Tests