fixed code duplication issue
added unit test