Fix tests
Merge pull request #7 from infinityb/master

Fix tests