Improve test coverage & fixes

Fixes #4