Refactor code & improve test coverage

Closes #1