Fix #336. Add unit test
Merge pull request #341 from palantir/fix-336

Fix #336. Add unit test