Tidy up pull request #7 and add tests, fixes #6