More tests

Fixed code and added more tests
Merge pull request #43 from alokin8/master

Looks good to me. (More tests)