Refactoring and added testing.

:heart: @Strnadj