Adding unit test and refactoring class.