added a unit test, and small refactoring