minor refactoring
added unit tests