Refactor and fix tests
Merge pull request #47 from jonmagic/fix-broken-integration-tests

Refactor and fix tests