add test and refactoring.
Merge pull request #4 from FumblePerson/bugfix

add test and refactoring.