Fix broken tests
Merge pull request #24 from square/edenman/fix-broken-test

Fix broken tests