Slight refactor. Added unit tests.